fix: local listener config

tinfoilanalytics · Nov 19, 2024 · 43e0cfc · 43e0cfc
1 parent 04a2140
commit 43e0cfc
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/examples/nginx/Dockerfile b/examples/nginx/Dockerfile
@@ -4,6 +4,6 @@ FROM nginxdemos/hello
 COPY --from=shim /nitro-attestation-shim /nitro-attestation-shim
 
 ENV NITRO_SHIM_PORT=6000
-ENV NITRO_UPSTREAM_PORT=80
+ENV NITRO_SHIM_UPSTREAM_PORT=80
 
 ENTRYPOINT ["/nitro-attestation-shim", "nginx", "-g", "daemon off;"]
diff --git a/examples/ollama/Dockerfile b/examples/ollama/Dockerfile
@@ -5,7 +5,7 @@ FROM ollama/ollama
 COPY --from=shim /nitro-attestation-shim /nitro-attestation-shim
 
 ENV NITRO_SHIM_PORT=6000
-ENV NITRO_UPSTREAM_PORT=11434
+ENV NITRO_SHIM_UPSTREAM_PORT=11434
 
 RUN nohup bash -c "ollama serve &" && sleep 5 && ollama pull llama3.2:1b
 

diff --git a/main.go b/main.go
@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -47,16 +48,23 @@ func getInt(env string, defaultValue int) int {
 
 func main() {
 	listenPort := getInt("NITRO_SHIM_PORT", 6000)
-	upstreamHost := fmt.Sprintf("localhost:%d", getInt("NITRO_UPSTREAM_PORT", 6001))
-	cid := getInt("NITRO_CID", 16)
+	upstreamHost := fmt.Sprintf("localhost:%d", getInt("NITRO_SHIM_UPSTREAM_PORT", 6001))
+	cid := getInt("NITRO_SHIM_CID", 16)
+	useVsock := os.Getenv("NITRO_SHIM_LOCAL") == ""
 
 	if len(os.Args) < 2 {
 		log.Fatalf("Usage: %s <command> [args...]", os.Args[0])
 	}
 
-	l, err := vsock.ListenContextID(uint32(cid), uint32(listenPort), nil)
+	var l net.Listener
+	var err error
+	if useVsock {
+		l, err = vsock.ListenContextID(uint32(cid), uint32(listenPort), nil)
+	} else {
+		l, err = net.Listen("tcp", fmt.Sprintf(":%d", listenPort))
+	}
 	if err != nil {
-		log.Fatalf("failed vsock.Listen: %s", err)
+		log.Fatalf("listen: %s", err)
 		return
 	}
 	defer l.Close()