package service import ( "context" "io" "net/http" "net/http/httptest" "strings" "testing" "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) func newOpenCodeSessionTestContext(t *testing.T, value string) *gin.Context { t.Helper() recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) if value != "" { c.Request.Header.Set(openCodeSessionHeader, value) } return c } func openCodeSessionTestService() *OpenAIGatewayService { return &OpenAIGatewayService{cfg: &config.Config{ Security: config.SecurityConfig{ URLAllowlist: config.URLAllowlistConfig{Enabled: false}, }, }} } func openCodeSessionTestAccount(baseURL string) *Account { return &Account{ ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{ "base_url": baseURL, credKeyHeaderOverrideEnabled: true, credKeyHeaderOverrides: map[string]any{"x-opencode-session": "fixed-account-value"}, }, } } func requireSingleOpenCodeSessionHeader(t *testing.T, headers http.Header, want string) { t.Helper() count := 0 for key, values := range headers { if strings.EqualFold(key, openCodeSessionHeader) { count += len(values) require.Equal(t, []string{want}, values) } } require.Equal(t, 1, count) } func TestApplyOpenCodeSessionHeaderTrustBoundary(t *testing.T) { tests := []struct { name string account *Account targetURL string incoming string want string }{ { name: "official origin", account: openCodeSessionTestAccount("https://opencode.ai/zen/v1"), targetURL: "https://opencode.ai/zen/v1/chat/completions", incoming: " conversation-123 ", want: "conversation-123", }, { name: "lookalike origin", account: openCodeSessionTestAccount("https://opencode.ai.evil.example/v1"), targetURL: "https://opencode.ai.evil.example/v1/responses", incoming: "conversation-123", }, { name: "subdomain is not implicitly trusted", account: openCodeSessionTestAccount("https://api.opencode.ai/v1"), targetURL: "https://api.opencode.ai/v1/responses", incoming: "conversation-123", }, { name: "insecure official origin", account: openCodeSessionTestAccount("http://opencode.ai/zen/v1"), targetURL: "http://opencode.ai/zen/v1/responses", incoming: "conversation-123", }, { name: "missing caller value", account: openCodeSessionTestAccount("https://opencode.ai/zen/v1"), targetURL: "https://opencode.ai/zen/v1/responses", }, { name: "oauth account", account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}, targetURL: "https://opencode.ai/zen/v1/responses", incoming: "conversation-123", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { headers := make(http.Header) applyOpenCodeSessionHeader(newOpenCodeSessionTestContext(t, tt.incoming), tt.account, tt.targetURL, headers) require.Equal(t, tt.want, headers.Get(openCodeSessionHeader)) }) } } func TestOpenCodeSessionForwardedByResponsesBuildersAfterAccountOverride(t *testing.T) { gin.SetMode(gin.TestMode) svc := openCodeSessionTestService() account := openCodeSessionTestAccount("https://opencode.ai/zen/v1") body := []byte(`{"model":"gpt-5","input":"hello"}`) tests := []struct { name string build func(*gin.Context) (*http.Request, error) }{ { name: "normal responses", build: func(c *gin.Context) (*http.Request, error) { return svc.buildUpstreamRequest(context.Background(), c, account, body, "token", false, "", false) }, }, { name: "passthrough responses", build: func(c *gin.Context) (*http.Request, error) { return svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), c, account, body, "token") }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { c := newOpenCodeSessionTestContext(t, "conversation-456") req, err := tt.build(c) require.NoError(t, err) requireSingleOpenCodeSessionHeader(t, req.Header, "conversation-456") }) } } func TestOpenCodeSessionMissingCallerValueKeepsExistingOverrideBehavior(t *testing.T) { gin.SetMode(gin.TestMode) svc := openCodeSessionTestService() account := openCodeSessionTestAccount("https://opencode.ai/zen/v1") c := newOpenCodeSessionTestContext(t, "") req, err := svc.buildUpstreamRequest( context.Background(), c, account, []byte(`{"model":"gpt-5","input":"hello"}`), "token", false, "", false, ) require.NoError(t, err) require.Equal(t, "fixed-account-value", getHeaderRaw(req.Header, "x-opencode-session")) } type openCodeSessionHTTPUpstream struct { request *http.Request } func (u *openCodeSessionHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { u.request = req return &http.Response{ StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(`{}`)), }, nil } func (u *openCodeSessionHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { return u.Do(req, proxyURL, accountID, concurrency) } func TestOpenCodeSessionForwardedByRawChatCompletionsAfterAccountOverride(t *testing.T) { gin.SetMode(gin.TestMode) upstream := &openCodeSessionHTTPUpstream{} svc := openCodeSessionTestService() svc.httpUpstream = upstream account := openCodeSessionTestAccount("https://opencode.ai/zen/v1") c := newOpenCodeSessionTestContext(t, "conversation-789") resp, err := svc.sendCCUpstreamRequest( context.Background(), c, account, "https://opencode.ai/zen/v1/chat/completions", []byte(`{"model":"gpt-5"}`), false, "token", "", "", ) require.NoError(t, err) require.NoError(t, resp.Body.Close()) require.NotNil(t, upstream.request) requireSingleOpenCodeSessionHeader(t, upstream.request.Header, "conversation-789") } func TestOpenCodeSessionIsNotForwardedToOtherUpstreams(t *testing.T) { gin.SetMode(gin.TestMode) svc := openCodeSessionTestService() body := []byte(`{"model":"gpt-5","input":"hello"}`) for _, baseURL := range []string{ "https://api.openai.com/v1", "https://opencode.ai.evil.example/v1", "https://api.opencode.ai/v1", } { t.Run(baseURL, func(t *testing.T) { account := &Account{ ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: map[string]any{"base_url": baseURL}, } c := newOpenCodeSessionTestContext(t, "private-conversation") req, err := svc.buildUpstreamRequest(context.Background(), c, account, body, "token", false, "", false) require.NoError(t, err) require.Empty(t, req.Header.Get(openCodeSessionHeader)) }) } }