55 "context"
66 "crypto/sha256"
77 "encoding/hex"
8+ "encoding/json"
89 "fmt"
910 "io"
1011 "net/http"
@@ -71,6 +72,7 @@ type CoalescingCache struct {
7172type ProxyService struct {
7273 config * Config
7374 httpClient * http.Client
75+ authService * AuthService
7476 workerPool WorkerPoolInterface
7577 circuitBreaker * CircuitBreaker
7678 bufferPool * sync.Pool
@@ -139,7 +141,7 @@ func (cc *CoalescingCache) CoalesceRequest(key string, fn func() interface{}) in
139141}
140142
141143// NewProxyService creates a new proxy service
142- func NewProxyService (cfg * Config , httpClient * http.Client , workerPool WorkerPoolInterface ) * ProxyService {
144+ func NewProxyService (cfg * Config , httpClient * http.Client , authService * AuthService , workerPool WorkerPoolInterface ) * ProxyService {
143145 circuitBreaker := & CircuitBreaker {
144146 state : CircuitClosed ,
145147 timeout : time .Duration (cfg .Timeouts .CircuitBreaker ) * time .Second ,
@@ -154,6 +156,7 @@ func NewProxyService(cfg *Config, httpClient *http.Client, workerPool WorkerPool
154156 return & ProxyService {
155157 config : cfg ,
156158 httpClient : httpClient ,
159+ authService : authService ,
157160 workerPool : workerPool ,
158161 circuitBreaker : circuitBreaker ,
159162 bufferPool : bufferPool ,
@@ -203,7 +206,18 @@ func (s *ProxyService) Handler() http.HandlerFunc {
203206 Error ("Worker error" , "error" , err )
204207 // Only write error if headers haven't been sent
205208 if ! respWrapper .headersSent {
206- http .Error (w , err .Error (), http .StatusInternalServerError )
209+ switch {
210+ case strings .Contains (err .Error (), "authentication error" ):
211+ http .Error (w , err .Error (), http .StatusUnauthorized )
212+ case strings .Contains (err .Error (), "token validation failed" ):
213+ http .Error (w , err .Error (), http .StatusUnauthorized )
214+ case strings .Contains (err .Error (), "bad request" ):
215+ http .Error (w , err .Error (), http .StatusBadRequest )
216+ case strings .Contains (err .Error (), "method not allowed" ):
217+ http .Error (w , err .Error (), http .StatusMethodNotAllowed )
218+ default :
219+ http .Error (w , err .Error (), http .StatusInternalServerError )
220+ }
207221 }
208222 }
209223 case <- ctx .Done ():
@@ -279,14 +293,40 @@ func (cb *CircuitBreaker) onFailure() {
279293func (s * ProxyService ) processProxyRequest (ctx context.Context , w http.ResponseWriter , r * http.Request ) error {
280294 Debug ("Starting proxy request" , "method" , r .Method , "path" , r .URL .Path )
281295
296+ // Validate method
297+ if r .Method != http .MethodPost {
298+ return fmt .Errorf ("method not allowed: %s" , r .Method )
299+ }
300+
282301 // Read the request body
283302 body , err := io .ReadAll (r .Body )
284303 if err != nil {
285304 Error ("Error reading request body" , "error" , err )
286- return NewProxyError ("read_request_body" , "failed to read request body" , err )
305+ // Check for "http: request body too large" error and return 413
306+ if strings .Contains (err .Error (), "http: request body too large" ) {
307+ return fmt .Errorf ("payload too large: %w" , err )
308+ }
309+ return fmt .Errorf ("bad request: failed to read request body: %w" , err )
287310 }
288311 defer r .Body .Close ()
289312
313+ // Basic body validation (for demonstration: consider empty body an error)
314+ if len (body ) == 0 {
315+ return fmt .Errorf ("bad request: empty request body" )
316+ }
317+
318+ // Strict JSON validation before authentication
319+ var js json.RawMessage
320+ if jsonErr := json .Unmarshal (body , & js ); jsonErr != nil {
321+ return fmt .Errorf ("bad request: invalid JSON: %w" , jsonErr )
322+ }
323+
324+ // Ensure we have a valid token before making the request
325+ if tokenErr := s .authService .EnsureValidToken (s .config ); tokenErr != nil {
326+ Error ("Failed to ensure valid token" , "error" , tokenErr )
327+ return NewAuthError ("token validation failed" , tokenErr )
328+ }
329+
290330 // Create new request to GitHub Copilot
291331 targetURL := copilotAPIBase + chatCompletionsPath
292332 Debug ("Sending request to target" , "url" , targetURL , "body_length" , len (body ))
0 commit comments