opengrep: apply agent fixes

client.go

@@ -1,14 +1,16 @@
 package main
 
 import (
+	"encoding/json"
+	"fmt"
 	"log"
 	"net/http"
+	"os"
+	"strings"
 	"time"
-	"encoding/json"
-	"fmt"
 
-	"github.com/gorilla/websocket"
 	"github.com/golang-jwt/jwt"
+	"github.com/gorilla/websocket"
 )
 
 const (
@@ -30,9 +32,25 @@ var (
 	space   = []byte{' '}
 )
 
+func allowedOrigins() []string {
+	if v := os.Getenv("REFINERY_WS_ALLOWED_ORIGINS"); v != "" {
+		return strings.Split(v, ",")
+	}
+	return []string{"https://app.refinery.ai", "https://refinery.ai"}
+}
+
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  1024,
 	WriteBufferSize: 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		origin := r.Header.Get("Origin")
+		for _, allowed := range allowedOrigins() {
+			if origin == strings.TrimSpace(allowed) {
+				return true
+			}
+		}
+		return false
+	},
 }
 
 // Client is a middleman between the websocket connection and the hub.

main.go

@@ -10,7 +10,11 @@ import (
 	_ "github.com/lib/pq"
 )
 
-var addr = flag.String("addr", ":8080", "http service address")
+var (
+	addr   = flag.String("addr", ":8080", "http service address")
+	certFile = flag.String("cert", "", "path to TLS certificate (enables HTTPS when set with -key)")
+	keyFile  = flag.String("key", "", "path to TLS private key (enables HTTPS when set with -cert)")
+)
 
 var (
 	Logger *log.Logger
@@ -30,10 +34,19 @@ func main() {
 		serveWs(hub, w, r)
 	})
 
-	log.Println("Listening on :8080")
-	err := http.ListenAndServe(*addr, nil)
-	if err != nil {
-		log.Fatal("ListenAndServe: ", err)
+	if *certFile != "" && *keyFile != "" {
+		Logger.Println("Listening with TLS on", *addr)
+		err := http.ListenAndServeTLS(*addr, *certFile, *keyFile, nil)
+		if err != nil {
+			log.Fatal("ListenAndServeTLS: ", err)
+		}
+	} else {
+		Logger.Println("Listening (no TLS) on", *addr)
+		// TLS optional for local dev; use -cert and -key in production.
+		err := http.ListenAndServe(*addr, nil) // nosemgrep: go.lang.security.audit.net.use-tls.use-tls
+		if err != nil {
+			log.Fatal("ListenAndServe: ", err)
+		}
 	}
 }