proxy.go

package main

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const freeSuffix = "-free"
const defaultOpenCodeAPIKey = "public"
const openCodeHost = "opencode.ai"
const openCodeUserAgent = "opencode/1.15.4 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.13"

var hopByHopHeaders = map[string]struct{}{
	"Connection":          {},
	"Keep-Alive":          {},
	"Proxy-Authenticate":  {},
	"Proxy-Authorization": {},
	"Te":                  {},
	"Trailer":             {},
	"Transfer-Encoding":   {},
	"Upgrade":             {},
}

type Proxy struct {
	cfg    Config
	client *http.Client
	logger *log.Logger
}

func NewProxy(cfg Config, logger *log.Logger) *Proxy {
	return &Proxy{
		cfg: cfg,
		client: &http.Client{
			Timeout: 0,
		},
		logger: logger,
	}
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/v1" && !strings.HasPrefix(r.URL.Path, "/v1/") {
		http.NotFound(w, r)
		return
	}

	if !p.authorized(r) {
		w.Header().Set("WWW-Authenticate", `Bearer realm="opencode-proxy"`)
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	if isModelsPath(r.URL.Path) {
		p.handleModels(w, r)
		return
	}

	p.forward(w, r)
}

func (p *Proxy) authorized(r *http.Request) bool {
	return r.Header.Get("Authorization") == "Bearer "+p.cfg.ProxyAPIKey
}

func isModelsPath(path string) bool {
	return strings.TrimRight(path, "/") == "/v1/models"
}

func (p *Proxy) handleModels(w http.ResponseWriter, r *http.Request) {
	resp, err := p.upstream(r, r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	if !p.cfg.IsFree || resp.StatusCode < 200 || resp.StatusCode >= 300 {
		copyResponse(w, resp, p.logger)
		return
	}

	body, err := readResponseBody(resp)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}

	filtered, ok := filterFreeModels(body)
	if !ok {
		copyHeader(w.Header(), resp.Header)
		w.WriteHeader(resp.StatusCode)
		_, _ = w.Write(body)
		return
	}

	copyHeader(w.Header(), resp.Header)
	w.Header().Del("Content-Length")
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)
	_, _ = w.Write(filtered)
}

func (p *Proxy) forward(w http.ResponseWriter, r *http.Request) {
	body := io.Reader(r.Body)
	if p.cfg.IsFree && r.Body != nil {
		newBody, err := appendFreeSuffixToRequestModel(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		body = newBody
	}

	resp, err := p.upstream(r, body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	copyResponse(w, resp, p.logger)
}

func (p *Proxy) upstream(r *http.Request, body io.Reader) (*http.Response, error) {
	upstreamURL := p.upstreamURL(r.URL)
	req, err := http.NewRequestWithContext(r.Context(), r.Method, upstreamURL, body)
	if err != nil {
		return nil, err
	}

	copyRequestHeaders(req.Header, r.Header)
	applyOpenCodeHeaders(req, p.cfg.OpenCodeAPIKey)
	req.Header.Set("X-Forwarded-Host", r.Host)
	appendForwardedFor(req.Header, r.RemoteAddr)

	return p.client.Do(req)
}

func applyOpenCodeHeaders(req *http.Request, apiKey string) {
	if apiKey == "" {
		apiKey = defaultOpenCodeAPIKey
	}

	req.Host = openCodeHost
	req.Header.Set("Authorization", "Bearer "+apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("User-Agent", openCodeUserAgent)
	req.Header.Set("X-Opencode-Client", "cli")
	req.Header.Set("X-Opencode-Project", "global")
	req.Header.Set("X-Opencode-Request", "msg_1")
	req.Header.Set("X-Opencode-Session", "ses_1")
	req.Header.Set("Connection", "keep-alive")
	req.Header.Set("Accept", "*/*")
	req.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd")
}

func (p *Proxy) upstreamURL(requestURL *url.URL) string {
	u := *p.cfg.OpenCodeBaseURL
	suffix := strings.TrimPrefix(requestURL.EscapedPath(), "/v1")
	if suffix == "" {
		suffix = "/"
	}
	u.Path = strings.TrimRight(u.Path, "/") + "/" + strings.TrimLeft(suffix, "/")
	u.RawQuery = requestURL.RawQuery
	return u.String()
}

func appendFreeSuffixToRequestModel(body io.ReadCloser) (io.Reader, error) {
	raw, err := io.ReadAll(body)
	if err != nil {
		return nil, err
	}
	_ = body.Close()
	if len(bytes.TrimSpace(raw)) == 0 {
		return bytes.NewReader(raw), nil
	}

	var payload map[string]any
	if err := json.Unmarshal(raw, &payload); err != nil {
		return bytes.NewReader(raw), nil
	}

	model, ok := payload["model"].(string)
	if !ok || model == "" || strings.HasSuffix(model, freeSuffix) {
		return bytes.NewReader(raw), nil
	}

	payload["model"] = model + freeSuffix
	rewritten, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}
	return bytes.NewReader(rewritten), nil
}

func filterFreeModels(body []byte) ([]byte, bool) {
	var payload map[string]any
	if err := json.Unmarshal(body, &payload); err != nil {
		return nil, false
	}

	data, ok := payload["data"].([]any)
	if !ok {
		return nil, false
	}

	filtered := make([]any, 0, len(data))
	for _, item := range data {
		model, ok := item.(map[string]any)
		if !ok {
			continue
		}
		id, ok := model["id"].(string)
		if !ok || !strings.HasSuffix(id, freeSuffix) {
			continue
		}

		copied := make(map[string]any, len(model))
		for key, value := range model {
			copied[key] = value
		}
		copied["id"] = strings.TrimSuffix(id, freeSuffix)
		filtered = append(filtered, copied)
	}

	payload["data"] = filtered
	rewritten, err := json.Marshal(payload)
	if err != nil {
		return nil, false
	}
	return rewritten, true
}

func readResponseBody(resp *http.Response) ([]byte, error) {
	reader, closeReader, err := responseBodyReader(resp)
	if err != nil {
		return nil, err
	}
	defer closeReader()

	return io.ReadAll(reader)
}

func copyResponse(w http.ResponseWriter, resp *http.Response, logger *log.Logger) {
	body, closeBody, err := responseBodyReader(resp)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer closeBody()

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)

	if flusher, ok := w.(http.Flusher); ok {
		buf := make([]byte, 32*1024)
		for {
			n, err := body.Read(buf)
			if n > 0 {
				_, _ = w.Write(buf[:n])
				flusher.Flush()
			}
			if err != nil {
				if err != io.EOF && logger != nil {
					logger.Printf("copy response: %v", err)
				}
				return
			}
		}
	}

	if _, err := io.Copy(w, body); err != nil && logger != nil {
		logger.Printf("copy response: %v", err)
	}
}

func responseBodyReader(resp *http.Response) (io.Reader, func(), error) {
	if !strings.EqualFold(resp.Header.Get("Content-Encoding"), "gzip") {
		return resp.Body, func() {}, nil
	}

	reader, err := gzip.NewReader(resp.Body)
	if err != nil {
		return nil, nil, err
	}

	resp.Header.Del("Content-Encoding")
	resp.Header.Del("Content-Length")
	return reader, func() { _ = reader.Close() }, nil
}

func copyHeader(dst, src http.Header) {
	for key, values := range src {
		if isHopByHopHeader(key) {
			continue
		}
		for _, value := range values {
			dst.Add(key, value)
		}
	}
}

func copyRequestHeaders(dst, src http.Header) {
	for key, values := range src {
		if isHopByHopHeader(key) ||
			strings.EqualFold(key, "Authorization") ||
			strings.EqualFold(key, "Accept-Encoding") {
			continue
		}
		for _, value := range values {
			dst.Add(key, value)
		}
	}
}

func isHopByHopHeader(key string) bool {
	_, ok := hopByHopHeaders[http.CanonicalHeaderKey(key)]
	return ok
}

func appendForwardedFor(header http.Header, remoteAddr string) {
	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		host = remoteAddr
	}
	if host == "" {
		return
	}

	if prior := header.Get("X-Forwarded-For"); prior != "" {
		header.Set("X-Forwarded-For", prior+", "+host)
		return
	}
	header.Set("X-Forwarded-For", host)
}

func newServer(cfg Config, logger *log.Logger) *http.Server {
	return &http.Server{
		Addr:              ":" + cfg.Port,
		Handler:           NewProxy(cfg, logger),
		ReadHeaderTimeout: 10 * time.Second,
	}
}