Auto retry on Cloudflare 400 / 524 errors

left-code.go

// Usage: go run left-code.go
// Configure: base_url = "http://127.0.0.1:1400/codex/v1"

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync/atomic"
	"time"
)

const (
	defaultListenAddr = "127.0.0.1:1400"
	defaultTargetBase = "https://www.right.codes"
	defaultTimeout    = 10 * time.Minute

	cfStatusTimeout = 524
)

var requestSeq uint64

func main() {
	listenAddr := flag.String("listen", defaultListenAddr, "listen address, e.g. 127.0.0.1:1400")
	targetBase := flag.String("target", defaultTargetBase, "target base URL, e.g. https://example.com")
	timeout := flag.Duration("timeout", defaultTimeout, "upstream timeout per request (0 = no timeout)")
	flag.Parse()

	log.SetOutput(os.Stderr)
	log.SetFlags(log.LstdFlags)

	targetBaseSanitized := strings.TrimRight(*targetBase, "/")
	log.Printf("Starting relay server on %s -> %s", *listenAddr, targetBaseSanitized)

	httpClient := newHTTPClient(*timeout)

	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		reqID := atomic.AddUint64(&requestSeq, 1)
		reqLogger := log.New(os.Stderr, fmt.Sprintf("[req=%d] ", reqID), log.LstdFlags|log.Lmsgprefix)

		urlString := r.URL.String()
		reqLogger.Printf("url=%s", urlString)

		statusCode := http.StatusOK
		defer func() {
			reqLogger.Printf("url=%s status_code=%d", urlString, statusCode)
		}()

		defer r.Body.Close()

		bodyBytes, err := io.ReadAll(r.Body)
		if err != nil {
			statusCode = http.StatusBadRequest
			reqLogger.Printf("Failed to read request body: %v", err)
			logRequest(reqLogger, r, nil)
			logResponseSummary(reqLogger, statusCode, nil)
			http.Error(w, "failed to read request body", statusCode)
			return
		}

		targetURL := targetBaseSanitized + r.URL.RequestURI()
		outboundReq, err := http.NewRequestWithContext(r.Context(), r.Method, targetURL, bytes.NewReader(bodyBytes))
		if err != nil {
			statusCode = http.StatusInternalServerError
			reqLogger.Printf("Failed to build outbound request: %v", err)
			logRequest(reqLogger, r, bodyBytes)
			logResponseSummary(reqLogger, statusCode, nil)
			http.Error(w, "failed to build request", statusCode)
			return
		}

		outboundReq.Header = r.Header.Clone()
		outboundReq.Host = ""

		writeHeaders := func(headers http.Header) {
			for key, values := range headers {
				for _, value := range values {
					if strings.EqualFold(key, "Transfer-Encoding") {
						continue
					}
					w.Header().Add(key, value)
				}
			}
		}

		streamResponse := func(resp *http.Response) {
			defer resp.Body.Close()

			statusCode = resp.StatusCode
			shouldLog := statusCode != http.StatusOK
			if shouldLog {
				logRequest(reqLogger, r, bodyBytes)
				logResponseSummary(reqLogger, statusCode, resp.Header)
			}

			writeHeaders(resp.Header)
			w.WriteHeader(statusCode)

			var loggedBody bytes.Buffer
			buf := make([]byte, 32*1024)

			for {
				n, readErr := resp.Body.Read(buf)
				if n > 0 {
					chunk := buf[:n]
					if shouldLog {
						loggedBody.Write(chunk)
					}
					if _, writeErr := w.Write(chunk); writeErr != nil {
						reqLogger.Printf("Failed to relay response body: %v", writeErr)
						return
					}
					if flusher, ok := w.(http.Flusher); ok {
						flusher.Flush()
					}
				}

				if readErr != nil {
					if readErr == io.EOF {
						break
					}
					reqLogger.Printf("Failed to read response body: %v", readErr)
					return
				}
			}

			if shouldLog {
				logBody(reqLogger, "Response Body", loggedBody.Bytes())
			}
		}

		writeBufferedResponse := func(status int, headers http.Header, body []byte) {
			statusCode = status

			shouldLog := statusCode != http.StatusOK
			if shouldLog {
				logRequest(reqLogger, r, bodyBytes)
				logResponseSummary(reqLogger, statusCode, headers)
				logBody(reqLogger, "Response Body", body)
			}

			writeHeaders(headers)
			w.WriteHeader(statusCode)
			if len(body) > 0 {
				if _, writeErr := w.Write(body); writeErr != nil {
					reqLogger.Printf("Failed to write response body: %v", writeErr)
					return
				}
				if flusher, ok := w.(http.Flusher); ok {
					flusher.Flush()
				}
			}
		}

		resp, err := httpClient.Do(outboundReq)
		if err != nil {
			statusCode = http.StatusBadGateway
			reqLogger.Printf("Relay request failed: %v", err)
			logRequest(reqLogger, r, bodyBytes)
			logResponseSummary(reqLogger, statusCode, nil)
			http.Error(w, "relay request failed", statusCode)
			return
		}

		if resp.StatusCode == http.StatusBadRequest || resp.StatusCode == cfStatusTimeout {
			firstBody, readErr := io.ReadAll(resp.Body)
			resp.Body.Close()
			if readErr != nil {
				statusCode = http.StatusBadGateway
				reqLogger.Printf("Failed to read response body: %v", readErr)
				logRequest(reqLogger, r, bodyBytes)
				logResponseSummary(reqLogger, statusCode, nil)
				http.Error(w, "failed to read response body", statusCode)
				return
			}

			tag := retryTag(resp.StatusCode, firstBody)
			if tag == "" {
				writeBufferedResponse(resp.StatusCode, resp.Header, firstBody)
				return
			}

			reqLogger.Printf("Meet %s for URL %s, will retry", tag, urlString)

			origStatus := resp.StatusCode
			origHeaders := resp.Header.Clone()
			origBody := firstBody

			backoff := 100 * time.Millisecond
			maxBackoff := 5 * time.Second

			var finalResp *http.Response
			finalStatus := origStatus
			finalHeaders := origHeaders
			finalBody := origBody

			for attempt := 1; attempt <= 10; attempt++ {
				select {
				case <-time.After(backoff):
				case <-r.Context().Done():
					return
				}

				retryLogger := log.New(os.Stderr, fmt.Sprintf("[req=%d.%d] ", reqID, attempt), log.LstdFlags|log.Lmsgprefix)

				retryReq, reqErr := http.NewRequestWithContext(r.Context(), r.Method, targetURL, bytes.NewReader(bodyBytes))
				if reqErr != nil {
					retryLogger.Printf("%s Retry #%d - Error: %v", tag, attempt, reqErr)
					backoff *= 2
					if backoff > maxBackoff {
						backoff = maxBackoff
					}
					continue
				}
				retryReq.Header = r.Header.Clone()
				retryReq.Host = ""

				resp2, doErr := httpClient.Do(retryReq)
				if doErr != nil {
					retryLogger.Printf("%s Retry #%d - Error: %v", tag, attempt, doErr)
					backoff *= 2
					if backoff > maxBackoff {
						backoff = maxBackoff
					}
					continue
				}

				retryLogger.Printf("%s Retry #%d - StatusCode %d", tag, attempt, resp2.StatusCode)

				if resp2.StatusCode != http.StatusBadRequest && resp2.StatusCode != cfStatusTimeout {
					finalResp = resp2
					break
				}

				body2, readErr := io.ReadAll(resp2.Body)
				resp2.Body.Close()
				if readErr != nil {
					retryLogger.Printf("%s Retry #%d - Error: %v", tag, attempt, readErr)
					backoff *= 2
					if backoff > maxBackoff {
						backoff = maxBackoff
					}
					continue
				}

				if retryTag(resp2.StatusCode, body2) == "" {
					finalStatus = resp2.StatusCode
					finalHeaders = resp2.Header.Clone()
					finalBody = body2
					break
				}

				backoff *= 2
				if backoff > maxBackoff {
					backoff = maxBackoff
				}
			}

			if finalResp != nil {
				streamResponse(finalResp)
				return
			}

			writeBufferedResponse(finalStatus, finalHeaders, finalBody)
			return
		}

		streamResponse(resp)
	})

	server := &http.Server{
		Addr:    *listenAddr,
		Handler: handler,
	}

	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("Server stopped with error: %v", err)
	}
}

func logRequest(logger *log.Logger, r *http.Request, body []byte) {
	logger.Printf("Request %s %s", r.Method, r.URL.String())
	logHeaders(logger, "Request Headers", r.Header)
	logBody(logger, "Request Body", body)
}

func logResponseSummary(logger *log.Logger, statusCode int, headers http.Header) {
	logger.Printf("Response %d %s", statusCode, http.StatusText(statusCode))
	logHeaders(logger, "Response Headers", headers)
}

func newHTTPClient(timeout time.Duration) *http.Client {
	transport, ok := http.DefaultTransport.(*http.Transport)
	if !ok {
		return &http.Client{Timeout: timeout}
	}

	cloned := transport.Clone()
	cloned.DialContext = (&net.Dialer{
		Timeout:   timeout,
		KeepAlive: 30 * time.Second,
	}).DialContext
	cloned.TLSHandshakeTimeout = timeout
	cloned.IdleConnTimeout = 0
	cloned.ResponseHeaderTimeout = 0
	cloned.ExpectContinueTimeout = 0

	client := &http.Client{Transport: cloned}
	if timeout > 0 {
		client.Timeout = timeout
	}
	return client
}

func logBody(logger *log.Logger, prefix string, body []byte) {
	if len(body) == 0 {
		logger.Printf("%s: <empty>", prefix)
		return
	}

	logger.Printf("%s:", prefix)
	for _, line := range strings.Split(string(body), "\n") {
		logger.Printf("%s", line)
	}
}

func logHeaders(logger *log.Logger, prefix string, headers http.Header) {
	if len(headers) == 0 {
		logger.Printf("%s: <none>", prefix)
		return
	}

	logger.Printf("%s:", prefix)
	for key, values := range headers {
		logger.Printf("%s: %s", key, strings.Join(values, ", "))
	}
}

func retryTag(statusCode int, body []byte) string {
	switch statusCode {
	case http.StatusBadRequest:
		if bytes.Contains(body, []byte("<title>HTTP Status 400 – Bad Request</title>")) {
			return "CF400"
		}
	case cfStatusTimeout:
		if bytes.Contains(body, []byte("524: A timeout occurred")) {
			return "CF524"
		}
	}
	return ""
}