Something went wrong. Try again.
small gleam coding and persistent agent daemon with a detachable cli and a canvas ui forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197// Package proxyserver forwards only the proxy extension's two public routes.package proxyserver
import ( "bytes" "context" "crypto/rand" "crypto/subtle" "encoding/hex" "encoding/json" "errors" "fmt" "io" "log" "net/http" "net/http/httputil" "os" "path/filepath" "time"
"albedo/cli/internal/config" "albedo/cli/internal/webui")
func KeyPath() string { return filepath.Join(config.HomeDir(), "proxy-key") }
// LoadKey creates or reads the home's private proxy key; rotate replaces it atomically.func LoadKey(rotate bool) (string, error) { path := KeyPath() if !rotate { key, err := readKey() if err == nil { return key, os.Chmod(path, 0o600) } if !errors.Is(err, os.ErrNotExist) { return "", err } } raw := make([]byte, 24) if _, err := rand.Read(raw); err != nil { return "", err } key := hex.EncodeToString(raw) if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { return "", err } if !rotate { f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) if errors.Is(err, os.ErrExist) { return readKey() } if err != nil { return "", err } _, err = f.WriteString(key) return key, errors.Join(err, f.Close()) } f, err := os.CreateTemp(filepath.Dir(path), ".proxy-key-*") if err != nil { return "", err } defer os.Remove(f.Name()) _, err = f.WriteString(key) if err = errors.Join(err, f.Close()); err != nil { return "", err } return key, os.Rename(f.Name(), path)}
func readKey() (string, error) { raw, err := os.ReadFile(KeyPath()) if err != nil { return "", err } key := string(raw) if _, err := hex.DecodeString(key); err != nil || len(key) != 48 { return "", fmt.Errorf("invalid proxy key in %s; use albedo proxy --new-key", KeyPath()) } return key, nil}
type Server struct { *webui.Access upstream webui.Upstream proxy *httputil.ReverseProxy logger *log.Logger}
func New(upstream webui.Upstream, output io.Writer) *Server { s := &Server{Access: webui.NewAccess(), upstream: upstream} if output != nil { s.logger = log.New(output, "", 0) } s.proxy = &httputil.ReverseProxy{ Rewrite: func(pr *httputil.ProxyRequest) { endpoint := s.upstream.Snapshot() pr.SetURL(webui.Loopback(endpoint.Port)) pr.Out.Host = pr.Out.URL.Host pr.Out.URL.Path = "/extensions/proxy" + pr.In.URL.Path pr.Out.URL.RawPath = "" for _, header := range []string{"Origin", "Cookie", "Referer", "Forwarded", "X-Forwarded-For", "X-Forwarded-Host", "X-Forwarded-Proto"} { pr.Out.Header.Del(header) } pr.Out.Header.Set("Authorization", "Bearer "+endpoint.Token) }, FlushInterval: -1, ModifyResponse: func(response *http.Response) error { response.Header.Del("Set-Cookie") response.Header.Del("Authorization") if response.StatusCode == http.StatusUnauthorized { return errors.New("daemon endpoint changed") } return nil }, ErrorHandler: func(w http.ResponseWriter, r *http.Request, _ error) { if r.Context().Err() != nil { return } s.refresh() w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusServiceUnavailable) _ = json.NewEncoder(w).Encode(map[string]map[string]string{"error": {"message": "the albedo daemon isn't running on this machine. start it by running albedo there, then retry.", "type": "service_unavailable"}}) }, } return s}
func (s *Server) refresh() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _ = s.upstream.Refresh(ctx)}
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { started := time.Now() response := &statusWriter{ResponseWriter: w, status: http.StatusOK} model := "-" if s.logger != nil { defer func() { s.logger.Printf("%q %d %s", model, response.status, time.Since(started).Round(time.Millisecond)) }() } if !s.Accept(r) { http.Error(response, "host, origin or network is not allowed", http.StatusForbidden) return } if r.URL.EscapedPath() != r.URL.Path || (r.URL.Path != "/v1/models" && r.URL.Path != "/v1/chat/completions") { http.NotFound(response, r) return } method := http.MethodGet if r.URL.Path == "/v1/chat/completions" { method = http.MethodPost } if r.Method != method { response.Header().Set("Allow", method) response.WriteHeader(http.StatusMethodNotAllowed) return } key, err := readKey() if err != nil || subtle.ConstantTimeCompare([]byte(r.Header.Get("Authorization")), []byte("Bearer "+key)) != 1 { response.Header().Set("WWW-Authenticate", "Bearer") http.Error(response, "proxy API key required", http.StatusUnauthorized) return } if r.Method == http.MethodPost && s.logger != nil { body, err := io.ReadAll(http.MaxBytesReader(response, r.Body, 32_000_000)) if err != nil { code := http.StatusBadRequest if _, oversized := errors.AsType[*http.MaxBytesError](err); oversized { code = http.StatusRequestEntityTooLarge } http.Error(response, "could not read request body", code) return } var request struct { Model string `json:"model"` } if json.Unmarshal(body, &request) == nil { model = request.Model[:min(len(request.Model), 256)] } r.Body = io.NopCloser(bytes.NewReader(body)) } s.proxy.ServeHTTP(response, r)}
type statusWriter struct { http.ResponseWriter status int}
func (w *statusWriter) WriteHeader(status int) { w.status = status w.ResponseWriter.WriteHeader(status)}func (w *statusWriter) Unwrap() http.ResponseWriter { return w.ResponseWriter }