diff --git a/Dockerfile b/Dockerfile index 2f144ad..0bfd925 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,4 +9,9 @@ COPY ${TARGETOS}/${TARGETARCH}/cockroachdb-mcp-server /usr/local/bin/cockroachdb USER nonroot:nonroot +# Distroless has no curl/shell, so the binary probes itself: GET /healthz in +# HTTP mode, no-op success in stdio mode. Orchestrator probes take precedence. +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \ + CMD ["/usr/local/bin/cockroachdb-mcp-server", "-healthcheck"] + ENTRYPOINT ["/usr/local/bin/cockroachdb-mcp-server"] diff --git a/README.md b/README.md index c73547d..e3905f7 100644 --- a/README.md +++ b/README.md @@ -123,6 +123,13 @@ export CRDB_MCP_ALLOW_INSECURE_HTTP=true > unauthenticated and `GET`/`HEAD`-only for orchestrator probes; all other > paths require the bearer token. +### Container health check + +The Docker image (distroless, no shell) defines a `HEALTHCHECK` that runs the +binary with `-healthcheck`: in HTTP mode it GETs `/healthz` on the configured +listen address (wildcard hosts rewritten to loopback); in stdio mode it exits +0. Orchestrator probes (k8s, ECS, Nomad) take precedence. + ### Tools shipped today | Tool | Description | diff --git a/main.go b/main.go index a2bc116..5c2db63 100644 --- a/main.go +++ b/main.go @@ -3,13 +3,16 @@ package main import ( "context" + "crypto/tls" stderrors "errors" "flag" "fmt" "io" + "net" "net/http" "os" "os/signal" + "strings" "syscall" "time" @@ -27,23 +30,35 @@ import ( ) const ( - serverName = "cockroachdb-mcp-server" - shutdownPeriod = 10 * time.Second - readHeaderTimeout = 10 * time.Second - readTimeout = 30 * time.Second - idleTimeout = 120 * time.Second - maxHeaderBytes = 32 << 10 // 32 KiB + serverName = "cockroachdb-mcp-server" + shutdownPeriod = 10 * time.Second + readHeaderTimeout = 10 * time.Second + readTimeout = 30 * time.Second + idleTimeout = 120 * time.Second + maxHeaderBytes = 32 << 10 // 32 KiB + healthCheckTimeout = 3 * time.Second ) var serverVersion = "0.1.0" func main() { showVersion := flag.Bool("version", false, "print version and exit") + healthCheck := flag.Bool("healthcheck", false, "probe local /healthz and exit 0/1; intended for Docker HEALTHCHECK in HTTP mode") flag.Parse() if *showVersion { fmt.Printf("%s %s\n", serverName, serverVersion) return } + if *healthCheck { + if !healthCheckEnabled() { + return + } + if err := probeHealthz(healthCheckURL()); err != nil { + fmt.Fprintf(os.Stderr, "healthcheck failed: %v\n", err) + os.Exit(1) + } + return + } // Bootstrap logger keeps config-load fatals visible. Stderr-only so a // bad CRDB_MCP_LOG_PATH can't swallow its own error. @@ -61,6 +76,53 @@ func main() { } } +// healthCheckEnabled reports whether the self-probe applies: /healthz only +// exists in HTTP mode, so stdio containers pass trivially. +func healthCheckEnabled() bool { + return strings.EqualFold(os.Getenv("CRDB_MCP_TRANSPORT"), config.TransportHTTP) +} + +// healthCheckURL derives the probe URL from the server's own listen env vars, +// rewriting wildcard hosts to loopback. +func healthCheckURL() string { + addr := os.Getenv("CRDB_MCP_HTTP_LISTEN_ADDR") + if addr == "" { + addr = ":8080" + } + if host, port, err := net.SplitHostPort(addr); err == nil { + switch host { + case "", "0.0.0.0", "::": + host = "127.0.0.1" + } + addr = net.JoinHostPort(host, port) + } + scheme := "http" + if os.Getenv("CRDB_MCP_TLS_CERT") != "" && os.Getenv("CRDB_MCP_TLS_KEY") != "" { + scheme = "https" + } + return scheme + "://" + addr + "/healthz" +} + +// probeHealthz issues a single GET and returns nil only on 200. TLS +// verification is skipped: this is a loopback self-check. +func probeHealthz(url string) error { + client := &http.Client{ + Timeout: healthCheckTimeout, + Transport: &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, //nolint:gosec // loopback self-probe + }, + } + resp, err := client.Get(url) + if err != nil { + return errors.Wrap(err, "probe healthz") + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return errors.Newf("/healthz returned %d", resp.StatusCode) + } + return nil +} + func run() error { cfg, err := config.Load() if err != nil { diff --git a/main_test.go b/main_test.go index 36ea769..f5010ff 100644 --- a/main_test.go +++ b/main_test.go @@ -75,6 +75,101 @@ func TestNewHTTPMux(t *testing.T) { }) } +func TestProbeHealthz(t *testing.T) { + t.Run("200 plain http returns nil", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + require.NoError(t, probeHealthz(srv.URL)) + }) + + t.Run("200 https with self-signed cert returns nil (skip-verify)", func(t *testing.T) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(srv.Close) + require.NoError(t, probeHealthz(srv.URL)) + }) + + t.Run("non-200 returns an error mentioning the status", func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + t.Cleanup(srv.Close) + err := probeHealthz(srv.URL) + require.Error(t, err) + require.Contains(t, err.Error(), "503") + }) + + t.Run("unreachable URL returns a wrapped error", func(t *testing.T) { + err := probeHealthz("http://127.0.0.1:1/healthz") + require.Error(t, err) + require.Contains(t, err.Error(), "probe healthz") + }) +} + +func TestHealthCheckURL(t *testing.T) { + t.Run("defaults to http://127.0.0.1:8080/healthz", func(t *testing.T) { + t.Setenv("CRDB_MCP_HTTP_LISTEN_ADDR", "") + t.Setenv("CRDB_MCP_TLS_CERT", "") + t.Setenv("CRDB_MCP_TLS_KEY", "") + require.Equal(t, "http://127.0.0.1:8080/healthz", healthCheckURL()) + }) + + t.Run("uses configured listen addr and rewrites :port to loopback", func(t *testing.T) { + t.Setenv("CRDB_MCP_HTTP_LISTEN_ADDR", ":9090") + t.Setenv("CRDB_MCP_TLS_CERT", "") + t.Setenv("CRDB_MCP_TLS_KEY", "") + require.Equal(t, "http://127.0.0.1:9090/healthz", healthCheckURL()) + }) + + t.Run("rewrites wildcard hosts to loopback", func(t *testing.T) { + t.Setenv("CRDB_MCP_TLS_CERT", "") + t.Setenv("CRDB_MCP_TLS_KEY", "") + for _, addr := range []string{"0.0.0.0:9090", "[::]:9090"} { + t.Setenv("CRDB_MCP_HTTP_LISTEN_ADDR", addr) + require.Equal(t, "http://127.0.0.1:9090/healthz", healthCheckURL()) + } + }) + + t.Run("preserves explicit non-wildcard host", func(t *testing.T) { + t.Setenv("CRDB_MCP_HTTP_LISTEN_ADDR", "10.0.0.5:9090") + t.Setenv("CRDB_MCP_TLS_CERT", "") + t.Setenv("CRDB_MCP_TLS_KEY", "") + require.Equal(t, "http://10.0.0.5:9090/healthz", healthCheckURL()) + }) + + t.Run("uses https when both TLS env vars are set", func(t *testing.T) { + t.Setenv("CRDB_MCP_HTTP_LISTEN_ADDR", ":8443") + t.Setenv("CRDB_MCP_TLS_CERT", "/etc/mcp/tls.crt") + t.Setenv("CRDB_MCP_TLS_KEY", "/etc/mcp/tls.key") + require.Equal(t, "https://127.0.0.1:8443/healthz", healthCheckURL()) + }) +} + +func TestHealthCheckEnabled(t *testing.T) { + t.Run("enabled in http mode", func(t *testing.T) { + t.Setenv("CRDB_MCP_TRANSPORT", "http") + require.True(t, healthCheckEnabled()) + }) + + t.Run("enabled regardless of case", func(t *testing.T) { + t.Setenv("CRDB_MCP_TRANSPORT", "HTTP") + require.True(t, healthCheckEnabled()) + }) + + t.Run("disabled in stdio mode", func(t *testing.T) { + t.Setenv("CRDB_MCP_TRANSPORT", "stdio") + require.False(t, healthCheckEnabled()) + }) + + t.Run("disabled when transport unset (stdio default)", func(t *testing.T) { + t.Setenv("CRDB_MCP_TRANSPORT", "") + require.False(t, healthCheckEnabled()) + }) +} + // TestRunHTTPShutdown verifies runHTTP returns cleanly when its context is // canceled, and that the listener is released so the port can be rebound. func TestRunHTTPShutdown(t *testing.T) {