fix: handle SEE contract

Author: williamleven

internal/app/backend.go

@@ -77,21 +77,26 @@ func NewBackend(ctx context.Context, rawCallbackURL string, signer Signer, deplo
 }
 
 type CloudEvent[Data any] struct {
-	Specversion string
-	Id          string
-	Source      string
-	Time        time.Time
-	Type        string
-	Data        Data
+	Specversion string    `json:"specversion"`
+	Id          string    `json:"id"`
+	Source      string    `json:"source"`
+	Type        string    `json:"type"`
+	Time        time.Time `json:"time"`
+	Data        Data      `json:"data"`
 }
 
 type ConnectBody struct {
 	ClientToken  string `json:"clientToken"`
 	SendToken    string `json:"sendToken"`
 	ConnectionId string `json:"connectionId"`
+	Origin       string `json:"origin"`
 }
 
-func (b *Backend) Connect(ctx context.Context, token string, signals chan<- Signal, maxLifeTime time.Duration) (string, error) {
+type DisconnectBody struct {
+	ConnectionId string `json:"connectionId"`
+}
+
+func (b *Backend) Connect(ctx context.Context, token string, origin string, signals chan<- Signal, maxLifeTime time.Duration) (string, error) {
 	connectionId := uuid.New().String()
 	sendToken, err := b.signer.CreateSendToken(SendTokenData{
 		ConnectionId: connectionId,
@@ -111,6 +116,7 @@ func (b *Backend) Connect(ctx context.Context, token string, signals chan<- Sign
 			ClientToken:  token,
 			SendToken:    string(sendToken),
 			ConnectionId: connectionId,
+			Origin:       origin,
 		},
 	}
 
@@ -153,13 +159,13 @@ func (b *Backend) Connect(ctx context.Context, token string, signals chan<- Sign
 func (b *Backend) Disconnect(ctx context.Context, connectionId string) error {
 	b.board.Unregister(connectionId)
 
-	event := CloudEvent[ConnectBody]{
+	event := CloudEvent[DisconnectBody]{
 		Specversion: "1.0",
 		Id:          uuid.New().String(),
 		Source:      b.deploymentURL,
 		Time:        time.Now(),
 		Type:        "charge.disconnected.v1",
-		Data: ConnectBody{
+		Data: DisconnectBody{
 			ConnectionId: connectionId,
 		},
 	}
@@ -181,10 +187,15 @@ func (b *Backend) Disconnect(ctx context.Context, connectionId string) error {
 	req.Header.Add("Content-Type", "application/cloudevents+json")
 	req.Header.Add("Webhook-Signature", string(signature))
 
-	_, err = http.DefaultClient.Do(req)
+	resp, err := http.DefaultClient.Do(req)
 	if err != nil {
 		return fmt.Errorf("post disconnected: %w", err)
 	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("disconnect callback returned %d", resp.StatusCode)
+	}
 
 	return nil
 }

internal/app/bouncer.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/url"
 	"slices"
 	"time"
 
@@ -24,6 +25,7 @@ type Bouncer struct {
 	m             *util.SyncMap[string, BounceStatus]
 	deploymentURL string
 	allowAll      bool
+	allowInsecure bool
 }
 
 type BounceStatus struct {
@@ -51,10 +53,11 @@ func (e NotAllowedError) Error() string {
 	return fmt.Sprintf("not allowed: %s (may try again after %s)", e.Reason, e.MayTryAgainAfter.Format(time.RFC3339))
 }
 
-func (b *Bouncer) Allowed(domain string) error {
+func (b *Bouncer) Allowed(callbackUrl *url.URL) error {
 	if b.allowAll {
 		return nil
 	}
+	domain := callbackUrl.Hostname()
 	if status, ok := b.m.Load(domain); ok && time.Now().Before(status.validBefore) {
 		return status.Allowed()
 	}
@@ -65,7 +68,7 @@ func (b *Bouncer) Allowed(domain string) error {
 			return st, nil
 		}
 
-		newStatus := b.fetchStatus(domain)
+		newStatus := b.fetchStatus(callbackUrl)
 		b.m.Store(domain, newStatus)
 		return newStatus, nil
 	})
@@ -76,12 +79,17 @@ func (b *Bouncer) Allowed(domain string) error {
 	return status.Allowed()
 }
 
-func (b *Bouncer) fetchStatus(domain string) BounceStatus {
+func (b *Bouncer) fetchStatus(callbackUrl *url.URL) BounceStatus {
 	status := BounceStatus{
 		validBefore: time.Now().Add(defaultNotAllowedCacheDuration),
 	}
 
-	allowedURL := fmt.Sprintf("https://%s/%s", domain, allowedInfoPath)
+	if !b.allowInsecure && callbackUrl.Scheme != "https" {
+		status.reason = "insecure scheme"
+		return status
+	}
+
+	allowedURL := fmt.Sprintf("%s://%s/%s", callbackUrl.Scheme, callbackUrl.Host, allowedInfoPath)
 	req, err := http.NewRequest(http.MethodGet, allowedURL, nil)
 	if err != nil {
 		status.reason = fmt.Sprintf("new request: %s", err)

internal/app/config.go

@@ -18,6 +18,7 @@ type Config struct {
 	RawSigningKeys        []byte        `envconfig:"SIGNING_KEYS" required:"true"`
 	MaxConnectionDuration time.Duration `envconfig:"MAX_CONNECTION_DURATION" default:"4h"`
 	AllowAllOrigins       bool          `envconfig:"ALLOW_ALL_ORIGINS"`
+	AllowInsecureOrigins  bool          `envconfig:"ALLOW_INSECURE_ORIGINS"`
 }
 
 type SigningKey struct {

internal/app/run.go

@@ -34,6 +34,7 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 		m:             &util.SyncMap[string, BounceStatus]{},
 		deploymentURL: cfg.DeploymentURL,
 		allowAll:      cfg.AllowAllOrigins,
+		allowInsecure: cfg.AllowInsecureOrigins,
 	}
 
 	bi := BackendIndex{
@@ -46,6 +47,11 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 	mux.HandleFunc("/sse", func(w http.ResponseWriter, r *http.Request) {
 		ctx = r.Context()
 
+		origin := r.Header.Get("Origin")
+		if origin != "" {
+			w.Header().Set("Access-Control-Allow-Origin", origin)
+		}
+
 		callback := r.URL.Query().Get("callback_url")
 		if callback == "" {
 			w.WriteHeader(http.StatusBadRequest)
@@ -58,7 +64,7 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 			return
 		}
 
-		bounceErr := bouncer.Allowed(callbackURL.Hostname())
+		bounceErr := bouncer.Allowed(callbackURL)
 		if bounceErr != nil {
 			if nbErr, ok := bounceErr.(NotAllowedError); ok {
 				w.Header().Set("Retry-After", fmt.Sprintf("%d", int(time.Until(nbErr.MayTryAgainAfter).Seconds())))
@@ -92,6 +98,14 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 			return
 		}
 
+		// Set SSE headers
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.Header().Set("Cache-Control", "no-cache")
+		w.Header().Set("Connection", "keep-alive")
+		// Write an initial comment to establish the SSE stream
+		fmt.Fprint(w, ":\n\n")
+		flusher.Flush()
+
 		maxDurationReached := time.After(cfg.MaxConnectionDuration)
 
 		wg := &sync.WaitGroup{}
@@ -106,7 +120,8 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 					if !ok {
 						return
 					}
-					_, err := w.Write(s.Message)
+					// Format SSE event: "data: <payload>\n\n"
+					_, err := fmt.Fprintf(w, "data: %s\n\n", s.Message)
 					if err != nil {
 						select {
 						case s.Result <- fmt.Errorf("write: %w", err):
@@ -126,8 +141,9 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 			}
 		})
 
-		id, err := be.Connect(ctx, token, signals, cfg.MaxConnectionDuration)
+		id, err := be.Connect(ctx, token, origin, signals, cfg.MaxConnectionDuration)
 		if err != nil {
+			// TODO: this error path is invalid as the headers have already been sent.
 			slog.Error("Failed to connect to backend", "backend_callback_url", be.callbackUrl.String(), "error", err)
 			w.WriteHeader(http.StatusServiceUnavailable)
 			return
@@ -138,10 +154,6 @@ func Run(ctx context.Context, log *slog.Logger, cfg Config) error {
 			}
 		}()
 
-		w.Header().Set("Content-Type", "text/event-stream")
-		w.Header().Set("Cache-Control", "no-cache")
-		w.Header().Set("Connection", "keep-alive")
-
 		wg.Wait()
 	})