Add TLS certificate generation, management utils #582
Conversation
Summary of ChangesHello @ericcurtin, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the security posture of the Model Runner API server by introducing comprehensive TLS capabilities. It provides a robust framework for generating, managing, and utilizing self-signed certificates, ensuring secure communication channels. The changes seamlessly integrate TLS options into the existing CLI commands and update the core Model Runner context, making it straightforward for users to enable and configure secure connections for their deployments. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- In main.go the handling of tlsServerErrors causes the process to exit immediately when TLS is disabled (the goroutine sends a nil error which wins the select); instead, only include tlsServerErrors in the select when TLS is actually started, or block that channel until shutdown so the non-TLS server can run normally.
- DetectContext assumes modelRunnerHost has an http/https scheme when constructing rawTLSURLPrefix for ModelRunnerEngineKindMobyManual; if a user passes a bare host like
localhost:12444, url.Parse will fail withmissing protocol scheme, so it would be safer to normalize or default the scheme before building TLS URLs. - For Desktop contexts, useTLS can be enabled via environment, but TLSURLPrefix is derived from an http URL and TLSClient is never constructed, so URL() may return a non-TLS URL while Client() is non-TLS as well; consider either ignoring MODEL_RUNNER_TLS for Desktop or making the behavior clearly consistent (e.g., fail fast or fully support TLS).
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In main.go the handling of tlsServerErrors causes the process to exit immediately when TLS is disabled (the goroutine sends a nil error which wins the select); instead, only include tlsServerErrors in the select when TLS is actually started, or block that channel until shutdown so the non-TLS server can run normally.
- DetectContext assumes modelRunnerHost has an http/https scheme when constructing rawTLSURLPrefix for ModelRunnerEngineKindMobyManual; if a user passes a bare host like `localhost:12444`, url.Parse will fail with `missing protocol scheme`, so it would be safer to normalize or default the scheme before building TLS URLs.
- For Desktop contexts, useTLS can be enabled via environment, but TLSURLPrefix is derived from an http URL and TLSClient is never constructed, so URL() may return a non-TLS URL while Client() is non-TLS as well; consider either ignoring MODEL_RUNNER_TLS for Desktop or making the behavior clearly consistent (e.g., fail fast or fully support TLS).
## Individual Comments
### Comment 1
<location> `cmd/cli/pkg/standalone/containers.go:384` </location>
<code_context>
+ // Get the directory containing the certificates
</code_context>
<issue_to_address>
**issue (bug_risk):** Mounted TLS certs in same directory may not match expected server.crt/server.key filenames
When `certDir == keyDir`, you mount the whole directory but still point `MODEL_RUNNER_TLS_CERT`/`MODEL_RUNNER_TLS_KEY` to `<tlsCertContainerPath>/server.crt` and `/server.key`. If the mounted files don’t use those names, TLS will fail. Either mount the specific files under those fixed names, or keep their original names in the container and set the env vars to the actual paths instead of hardcoding `server.crt`/`server.key`.
</issue_to_address>
### Comment 2
<location> `cmd/cli/desktop/context.go:261-265` </location>
<code_context>
case types.ModelRunnerEngineKindCloud:
rawURLPrefix = "http://localhost:" + strconv.Itoa(standalone.DefaultControllerPortCloud)
+ rawTLSURLPrefix = "https://localhost:" + strconv.Itoa(standalone.DefaultTLSPortCloud)
case types.ModelRunnerEngineKindMobyManual:
rawURLPrefix = modelRunnerHost
+ // For manual hosts, derive TLS URL from the provided host
+ if useTLS && !strings.HasPrefix(modelRunnerHost, "https://") {
+ rawTLSURLPrefix = strings.Replace(modelRunnerHost, "http://", "https://", 1)
+ } else {
+ rawTLSURLPrefix = modelRunnerHost
</code_context>
<issue_to_address>
**suggestion (bug_risk):** TLS URL derivation for manual host can produce unexpected or empty scheme scenarios
For `ModelRunnerEngineKindMobyManual`, the TLS URL is built by string-replacing `http://` with `https://` when `useTLS` is true. If `modelRunnerHost` has no scheme (e.g. `localhost:12435`) or an unexpected scheme, `rawTLSURLPrefix` may be unchanged or non-HTTPS, and `url.Parse` of a bare `host:port` will be treated as a path, not a host. Consider normalizing the input by (1) ensuring a scheme is set (default to `https://` when `useTLS` is true and none is provided), and (2) validating that the final parsed `tlsURLPrefix` uses `https` when TLS is enabled.
Suggested implementation:
```golang
case types.ModelRunnerEngineKindMobyManual:
normalizedHost := modelRunnerHost
// Ensure the manual host has a scheme.
// Default to https when TLS is requested, otherwise http.
if !strings.HasPrefix(normalizedHost, "http://") && !strings.HasPrefix(normalizedHost, "https://") {
if useTLS {
normalizedHost = "https://" + normalizedHost
} else {
normalizedHost = "http://" + normalizedHost
}
}
rawURLPrefix = normalizedHost
// Derive TLS URL from the normalized host, ensuring https when TLS is enabled.
if useTLS {
if strings.HasPrefix(normalizedHost, "http://") {
rawTLSURLPrefix = "https://" + strings.TrimPrefix(normalizedHost, "http://")
} else {
rawTLSURLPrefix = normalizedHost
}
} else {
rawTLSURLPrefix = normalizedHost
}
```
To fully implement your suggestion about validating the TLS scheme:
1. In the section where `rawURLPrefix` and `rawTLSURLPrefix` are parsed (likely using `url.Parse`), add a check such as:
- After parsing `tlsURLPrefix`, if `useTLS` is true and `tlsURLPrefix.Scheme != "https"`, return an error indicating that TLS was requested but the URL is not HTTPS.
2. Ensure `net/url` is imported if it isn't already (it likely is, since `url.Parse` is mentioned).
</issue_to_address>
### Comment 3
<location> `pkg/tls/certs_test.go:276-285` </location>
<code_context>
+ }
+}
+
+func TestCertsExistAndValid(t *testing.T) {
+ // Create a temporary directory
+ tmpDir, err := os.MkdirTemp("", "certs-test-*")
+ if err != nil {
+ t.Fatalf("Failed to create temp dir: %v", err)
+ }
+ defer os.RemoveAll(tmpDir)
+
+ paths := &CertPaths{
+ CACert: filepath.Join(tmpDir, CACertFile),
+ CAKey: filepath.Join(tmpDir, CAKeyFile),
+ ServerCert: filepath.Join(tmpDir, ServerCertFile),
+ ServerKey: filepath.Join(tmpDir, ServerKeyFile),
+ }
+
+ // Should return false when files don't exist
+ if certsExistAndValid(paths) {
+ t.Error("certsExistAndValid() should return false for non-existent files")
+ }
+
+ // Generate certs
+ err = GenerateCertificates(paths)
+ if err != nil {
+ t.Fatalf("GenerateCertificates() error = %v", err)
+ }
+
+ // Should return true for valid certs
+ if !certsExistAndValid(paths) {
+ t.Error("certsExistAndValid() should return true for valid certs")
+ }
</code_context>
<issue_to_address>
**suggestion (testing):** Extend certsExistAndValid tests to cover the "near-expiry" case where certificates should be treated as invalid.
This test covers only the non-existent and valid-cert scenarios, but not the case where certificates expiring within 30 days are treated as invalid. Please add a subtest that writes a server cert with NotAfter around time.Now().Add(15 * 24 * time.Hour) to paths.ServerCert and asserts certsExistAndValid(paths) returns false, e.g., by creating the cert via x509.CreateCertificate and reusing existing helpers where possible.
Suggested implementation:
```golang
func TestCertsExistAndValid(t *testing.T) {
// Create a temporary directory
tmpDir, err := os.MkdirTemp("", "certs-test-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer os.RemoveAll(tmpDir)
paths := &CertPaths{
CACert: filepath.Join(tmpDir, CACertFile),
CAKey: filepath.Join(tmpDir, CAKeyFile),
ServerCert: filepath.Join(tmpDir, ServerCertFile),
ServerKey: filepath.Join(tmpDir, ServerKeyFile),
}
// Should return false when files don't exist
if certsExistAndValid(paths) {
t.Error("certsExistAndValid() should return false for non-existent files")
}
// Generate certs
err = GenerateCertificates(paths)
if err != nil {
t.Fatalf("GenerateCertificates() error = %v", err)
}
// Should return true for valid certs
if !certsExistAndValid(paths) {
t.Error("certsExistAndValid() should return true for valid certs")
}
t.Run("near-expiry certs treated as invalid", func(t *testing.T) {
// Load CA certificate
caCertPEM, err := os.ReadFile(paths.CACert)
if err != nil {
t.Fatalf("failed to read CA cert: %v", err)
}
caCertBlock, _ := pem.Decode(caCertPEM)
if caCertBlock == nil {
t.Fatal("failed to decode CA cert PEM")
}
caCert, err := x509.ParseCertificate(caCertBlock.Bytes)
if err != nil {
t.Fatalf("failed to parse CA cert: %v", err)
}
// Load CA private key
caKeyPEM, err := os.ReadFile(paths.CAKey)
if err != nil {
t.Fatalf("failed to read CA key: %v", err)
}
caKeyBlock, _ := pem.Decode(caKeyPEM)
if caKeyBlock == nil {
t.Fatal("failed to decode CA key PEM")
}
caKey, err := x509.ParsePKCS1PrivateKey(caKeyBlock.Bytes)
if err != nil {
t.Fatalf("failed to parse CA key: %v", err)
}
// Create a server certificate that expires in 15 days
serverKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("failed to generate server key: %v", err)
}
template := &x509.Certificate{
SerialNumber: big.NewInt(2),
NotBefore: time.Now().Add(-1 * time.Hour),
NotAfter: time.Now().Add(15 * 24 * time.Hour), // within 30 days, should be treated as invalid
KeyUsage: x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
DNSNames: []string{"localhost"},
}
derBytes, err := x509.CreateCertificate(rand.Reader, template, caCert, &serverKey.PublicKey, caKey)
if err != nil {
t.Fatalf("failed to create near-expiry server cert: %v", err)
}
// Overwrite server cert and key on disk with the near-expiry ones
serverCertPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
if err := os.WriteFile(paths.ServerCert, serverCertPEM, 0o600); err != nil {
t.Fatalf("failed to write near-expiry server cert: %v", err)
}
serverKeyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(serverKey)})
if err := os.WriteFile(paths.ServerKey, serverKeyPEM, 0o600); err != nil {
t.Fatalf("failed to write near-expiry server key: %v", err)
}
// Now certsExistAndValid should treat these near-expiry certs as invalid
if certsExistAndValid(paths) {
t.Error("certsExistAndValid() should return false for near-expiry certificates")
}
})
}
```
You will also need to ensure the following imports exist in `pkg/tls/certs_test.go` (merged with whatever is already there):
- `"crypto/rand"`
- `"crypto/rsa"`
- `"crypto/x509"`
- `"encoding/pem"`
- `"math/big"`
- `"time"`
`os` and `path/filepath` are already required by the existing `TestCertsExistAndValid`, so if they’re not imported yet, add those as well. Adjust the import block accordingly to avoid unused-import errors.
</issue_to_address>
### Comment 4
<location> `pkg/tls/certs_test.go:185-194` </location>
<code_context>
+func TestLoadClientTLSConfig(t *testing.T) {
</code_context>
<issue_to_address>
**suggestion (testing):** Add negative-path tests for LoadClientTLSConfig when CA loading fails.
Right now this only covers happy paths. Please add subtests for failing CA load, e.g. (1) non-existent CA path and (2) existing file with invalid PEM, and assert that LoadClientTLSConfig returns a non-nil error and doesn’t configure TLS in those cases.
Suggested implementation:
```golang
if len(tlsConfig.Certificates) != 1 {
t.Errorf("Expected 1 certificate, got %d", len(tlsConfig.Certificates))
}
if tlsConfig.MinVersion != tls.VersionTLS12 {
t.Error("TLS minimum version should be TLS 1.2")
}
}
func TestLoadClientTLSConfig_CAErrors(t *testing.T) {
tmpDir := t.TempDir()
t.Run("non-existent CA path", func(t *testing.T) {
paths := &CertPaths{
CACert: filepath.Join(tmpDir, "does-not-exist-ca.pem"),
}
tlsConfig, err := LoadClientTLSConfig(paths)
if err == nil {
t.Fatalf("expected error when loading client TLS config with non-existent CA path, got nil")
}
if tlsConfig != nil {
t.Fatalf("expected nil TLS config on CA load error, got non-nil")
}
})
t.Run("invalid CA PEM", func(t *testing.T) {
caPath := filepath.Join(tmpDir, "invalid-ca.pem")
if err := os.WriteFile(caPath, []byte("this is not a valid PEM"), 0o600); err != nil {
t.Fatalf("failed to write invalid CA file: %v", err)
}
paths := &CertPaths{
CACert: caPath,
}
tlsConfig, err := LoadClientTLSConfig(paths)
if err == nil {
t.Fatalf("expected error when loading client TLS config with invalid CA PEM, got nil")
}
if tlsConfig != nil {
t.Fatalf("expected nil TLS config on CA load error, got non-nil")
}
})
}
func TestLoadClientTLSConfig(t *testing.T) {
```
These changes assume:
1. `CertPaths` and `LoadClientTLSConfig` are in the same package as the tests and that `LoadClientTLSConfig` returns `(*tls.Config, error)` and returns `nil` for the config on error.
2. `os` and `path/filepath` are already imported in `pkg/tls/certs_test.go`. If they are not, you will need to add:
- `import "os"`
- `import "path/filepath"`
If `LoadClientTLSConfig` returns a non-nil `*tls.Config` even on error, adjust the `tlsConfig != nil` assertions accordingly (for example, asserting that it has no `RootCAs` or other security-relevant fields configured).
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive TLS support for the Model Runner API server, including automatic certificate generation and management. The changes span across CLI commands, container creation logic, and the main server application, adding new flags and environment variables for TLS configuration. The implementation is robust, with a new pkg/tls that follows security best practices and includes thorough testing. My review focuses on a few opportunities to reduce code duplication and improve maintainability by refactoring repeated logic into helper functions and using constants instead of magic strings. Overall, this is a solid contribution that significantly enhances the security of the model runner.
This introduces a new pkg/tls package that provides automatic TLS certificate generation for the Model Runner API server. The feature includes self-signed CA and server certificate creation with proper validation and renewal checks. Additionally, TLS support is integrated into the install-runner, reinstall-runner, and start-runner commands with new flags for enabling TLS, specifying ports, and providing custom certificate paths. The Model Runner context now supports TLS-enabled connections with proper client configuration. Signed-off-by: Eric Curtin <eric.curtin@docker.com>
doringeman
left a comment
doringeman
left a comment
There was a problem hiding this comment.
LGTM!
FWIW to test with DMR in Docker Desktop:
make docker-build DOCKER_IMAGE=docker/model-runner:test-tls
_MODEL_RUNNER_TREAT_DESKTOP_AS_MOBY=1 MODEL_RUNNER_CONTROLLER_VERSION=test-tls docker model start-runner --tls --tls-cert tls-cert
$ curl -k https://localhost:12444/v1/models
{"object":"list","data":[{"id":"ai/llama3.2:latest","object":"model","created":1742916473,"owned_by":"docker"},{"id":"ai/smollm2:latest","object":"model","created":1742816981,"owned_by":"docker"}]}
$ curl https://localhost:12444/v1/models
curl: (60) SSL certificate problem: unable to get local issuer certificate
| // GetCACertPath returns the path to the CA certificate file. | ||
| // Returns the custom path if provided, otherwise returns the default path. | ||
| func GetCACertPath(customPath string) (string, error) { | ||
| if customPath != "" { | ||
| return customPath, nil | ||
| } | ||
|
|
||
| paths, err := DefaultCertPaths() | ||
| if err != nil { | ||
| return "", err | ||
| } | ||
| return paths.CACert, nil | ||
| } |
There was a problem hiding this comment.
Dead code.
| // GetCACertPath returns the path to the CA certificate file. | |
| // Returns the custom path if provided, otherwise returns the default path. | |
| func GetCACertPath(customPath string) (string, error) { | |
| if customPath != "" { | |
| return customPath, nil | |
| } | |
| paths, err := DefaultCertPaths() | |
| if err != nil { | |
| return "", err | |
| } | |
| return paths.CACert, nil | |
| } |
| IsCA: true, | ||
| KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign | x509.KeyUsageDigitalSignature, | ||
| BasicConstraintsValid: true, | ||
| MaxPathLen: 1, |
There was a problem hiding this comment.
| MaxPathLen: 1, | |
| MaxPathLen: 0, | |
| MaxPathLenZero: true, |
?
Why do we allow an intermediate CA?
| NotAfter: time.Now().AddDate(0, 0, DefaultCertValidityDays), | ||
| KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment, | ||
| ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, | ||
| DNSNames: []string{"localhost", "docker-model-runner"}, |
There was a problem hiding this comment.
I wonder if we'll need to add another one here for Offload.
| if tlsOpts.Enabled { | ||
| env = append(env, "MODEL_RUNNER_TLS_ENABLED=true") | ||
| env = append(env, "MODEL_RUNNER_TLS_PORT="+strconv.Itoa(int(tlsPort))) | ||
| if tlsOpts.CertPath != "" && tlsOpts.KeyPath != "" { |
There was a problem hiding this comment.
It looks like both on the client side and on the server side, both certificate and key paths must be provided, or neither. But this is silently ignored and we even proceed with the auto-generated certs if only one is provided.
|
FWIW this fixes #445. |
|
@doringeman I'm gonna merge this to unblock you if you want to look at this futher, feel free to make any suggested changes, even breaking ones, I'll approve the PR. One thing that's important when we implement this, we have to retain the ability to do plain-text http also by default, because many clients/tools out there will assume http as if they were connection to a plain-text vllm, llama.cpp, Ollama, instance etc. |
|
And yeah I did some autogenerating self-signed stuff here, that's not great, but it's better than no encryption. It's still encrypted, not authenticated. Even with this, it's more secure than most other inference engines which do zero and maybe doing zero is not an unwise choice, in a production environment, maybe you do just put nginx in front of this or something similar to do tls encryption and forget about it. |
|
Somehow this broken the build on main, even though it passed here. A linter build |
This introduces a new pkg/tls package that provides automatic TLS
certificate generation for the Model Runner API server. The feature
includes self-signed CA and server certificate creation with proper
validation and renewal checks. Additionally, TLS support is integrated
into the install-runner, reinstall-runner, and start-runner commands
with new flags for enabling TLS, specifying ports, and providing custom
certificate paths. The Model Runner context now supports TLS-enabled
connections with proper client configuration.
Signed-off-by: Eric Curtin eric.curtin@docker.com