diff --git a/agent.go b/agent.go index 3893ded..633d76f 100644 --- a/agent.go +++ b/agent.go @@ -318,7 +318,7 @@ func handleMessage( URL string `json:"url"` Checksum string `json:"checksum"` } - if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" { + if err := json.Unmarshal(msg.Payload, &payload); err != nil || payload.URL == "" || payload.Checksum == "" { slog.Error("invalid update payload") return } diff --git a/agent_test.go b/agent_test.go index de2d68c..12163f5 100644 --- a/agent_test.go +++ b/agent_test.go @@ -239,11 +239,33 @@ func TestHandleMessage(t *testing.T) { mtCh := make(chan *pb.MikrotikResult, 1) credCh := make(chan *pb.CredentialTestResult, 1) monCh := make(chan *pb.MonitoringCheck, 1) - payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"}) + payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"}) handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) // Should log error but not panic }) + t.Run("update missing checksum", func(t *testing.T) { + origUpdate := doSelfUpdate + defer func() { doSelfUpdate = origUpdate }() + + called := false + doSelfUpdate = func(url, checksum string) error { + called = true + return nil + } + + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"}) + handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh) + + if called { + t.Error("selfUpdate should not be called with empty checksum") + } + }) + t.Run("unknown event", func(t *testing.T) { snmpCh := make(chan *pb.SnmpResult, 1) mtCh := make(chan *pb.MikrotikResult, 1) diff --git a/mikrotik.go b/mikrotik.go index 644c275..47f4da0 100644 --- a/mikrotik.go +++ b/mikrotik.go @@ -14,8 +14,9 @@ import ( ) const ( - mikrotikConnTimeout = 30 * time.Second - mikrotikReadTimeout = 30 * time.Second + mikrotikConnTimeout = 30 * time.Second + mikrotikReadTimeout = 30 * time.Second + maxMikrotikWordSize = 10 << 20 // 10 MB ) var mikrotikDial = mikrotikConnect @@ -170,6 +171,9 @@ func (c *mikrotikClient) readWord() (string, error) { if length == 0 { return "", nil } + if length > maxMikrotikWordSize { + return "", fmt.Errorf("word size %d exceeds max %d", length, maxMikrotikWordSize) + } buf := make([]byte, length) if _, err := io.ReadFull(c.conn, buf); err != nil { return "", fmt.Errorf("read word: %w", err) diff --git a/mikrotik_test.go b/mikrotik_test.go index 0b4c80d..bc8105e 100644 --- a/mikrotik_test.go +++ b/mikrotik_test.go @@ -567,6 +567,34 @@ func TestReadLength5ByteError(t *testing.T) { } } +func TestReadWordExceedsMaxSize(t *testing.T) { + oversize := maxMikrotikWordSize + 1 + var buf bytes.Buffer + buf.Write(encodeLength(oversize)) + // Don't need to write the payload — should reject before reading it + c := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + _, err := c.readWord() + if err == nil { + t.Error("expected error for word exceeding max size") + } + if !containsStr(err.Error(), "exceeds max") { + t.Errorf("expected 'exceeds max' in error, got: %v", err) + } +} + +func containsStr(s, sub string) bool { + return len(s) >= len(sub) && searchStr(s, sub) +} + +func searchStr(s, sub string) bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} + func TestReadSentenceError(t *testing.T) { // Buffer with valid length byte but truncated word data buf := bytes.NewBuffer([]byte{0x03, 'a'}) // length=3 but only 1 byte of data diff --git a/update.go b/update.go index cdca9af..77dcc38 100644 --- a/update.go +++ b/update.go @@ -2,10 +2,12 @@ package main import ( "crypto/sha256" + "crypto/subtle" "fmt" "io" "log/slog" "net/http" + "net/url" "os" "syscall" ) @@ -13,12 +15,25 @@ import ( var osExecutable = os.Executable var osWriteFile = os.WriteFile var osRename = os.Rename +var httpGet = http.Get +var maxUpdateSize int64 = 100 << 20 // 100 MB // selfUpdate downloads a new binary, verifies its checksum, replaces the current binary, and re-execs. func selfUpdate(downloadURL, expectedChecksum string) error { + u, err := url.Parse(downloadURL) + if err != nil { + return fmt.Errorf("parse url: %w", err) + } + if u.Scheme != "https" { + return fmt.Errorf("HTTPS required for update URL, got %q", u.Scheme) + } + if expectedChecksum == "" { + return fmt.Errorf("checksum required for update") + } + slog.Info("downloading update", "url", downloadURL) - resp, err := http.Get(downloadURL) + resp, err := httpGet(downloadURL) if err != nil { return fmt.Errorf("download: %w", err) } @@ -28,20 +43,21 @@ func selfUpdate(downloadURL, expectedChecksum string) error { return fmt.Errorf("download failed: status %d", resp.StatusCode) } - body, err := io.ReadAll(resp.Body) + body, err := io.ReadAll(io.LimitReader(resp.Body, maxUpdateSize+1)) if err != nil { return fmt.Errorf("read body: %w", err) } + if int64(len(body)) > maxUpdateSize { + return fmt.Errorf("download size %d exceeds max %d", len(body), maxUpdateSize) + } slog.Info("downloaded update", "bytes", len(body)) - // Verify SHA256 checksum - if expectedChecksum != "" { - actual := fmt.Sprintf("%x", sha256.Sum256(body)) - if actual != expectedChecksum { - return fmt.Errorf("checksum mismatch: expected %s, got %s", expectedChecksum, actual) - } - slog.Info("checksum verified") + // Verify SHA256 checksum (constant-time comparison) + actual := fmt.Sprintf("%x", sha256.Sum256(body)) + if subtle.ConstantTimeCompare([]byte(actual), []byte(expectedChecksum)) != 1 { + return fmt.Errorf("checksum mismatch: expected %s, got %s", expectedChecksum, actual) } + slog.Info("checksum verified") // Write to temp file next to current binary currentExe, err := osExecutable() diff --git a/update_test.go b/update_test.go index 9679e58..373ba11 100644 --- a/update_test.go +++ b/update_test.go @@ -10,32 +10,60 @@ import ( "testing" ) +func TestSelfUpdateRejectsHTTP(t *testing.T) { + err := selfUpdate("http://example.com/agent", "abc123") + if err == nil { + t.Error("expected error for HTTP URL") + } + if !strings.Contains(err.Error(), "HTTPS required") { + t.Errorf("expected 'HTTPS required' in error, got: %v", err) + } +} + +func TestSelfUpdateRequiresChecksum(t *testing.T) { + err := selfUpdate("https://example.com/agent", "") + if err == nil { + t.Error("expected error for empty checksum") + } + if !strings.Contains(err.Error(), "checksum required") { + t.Errorf("expected 'checksum required' in error, got: %v", err) + } +} + func TestSelfUpdateBadURL(t *testing.T) { - err := selfUpdate("http://127.0.0.1:1/nonexistent", "") + err := selfUpdate("https://127.0.0.1:1/nonexistent", "abc123") if err == nil { t.Error("expected error for unreachable URL") } } func TestSelfUpdateChecksumMismatch(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("fake binary")) })) defer srv.Close() - err := selfUpdate(srv.URL, "0000000000000000000000000000000000000000000000000000000000000000") + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL), "0000000000000000000000000000000000000000000000000000000000000000") if err == nil { t.Error("expected checksum mismatch error") } } func TestSelfUpdate404(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNotFound) })) defer srv.Close() - err := selfUpdate(srv.URL+"/missing", "") + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL)+"/missing", "abc123") if err == nil { t.Error("expected error for 404 response") } @@ -46,7 +74,7 @@ func TestSelfUpdate404(t *testing.T) { func TestSelfUpdateReadBodyError(t *testing.T) { // Server sends Content-Length header but closes connection prematurely - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Length", "99999") w.WriteHeader(http.StatusOK) _, _ = w.Write([]byte("partial")) @@ -57,7 +85,11 @@ func TestSelfUpdateReadBodyError(t *testing.T) { })) defer srv.Close() - err := selfUpdate(srv.URL, "") + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL), "abc123") // This may or may not error depending on io.ReadAll behavior with truncated body // but we exercise the code path _ = err @@ -70,12 +102,19 @@ func TestSelfUpdateOsExecutableError(t *testing.T) { return "", fmt.Errorf("executable not found") } - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte("binary data")) + body := []byte("binary data") + checksum := fmt.Sprintf("%x", sha256.Sum256(body)) + + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(body) })) defer srv.Close() - err := selfUpdate(srv.URL, "") + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) if err == nil { t.Error("expected os.Executable error") } @@ -96,12 +135,19 @@ func TestSelfUpdateWriteFileError(t *testing.T) { return fmt.Errorf("disk full") } - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte("binary data")) + body := []byte("binary data") + checksum := fmt.Sprintf("%x", sha256.Sum256(body)) + + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(body) })) defer srv.Close() - err := selfUpdate(srv.URL, "") + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) if err == nil { t.Error("expected write file error") } @@ -125,12 +171,19 @@ func TestSelfUpdateRenameError(t *testing.T) { return fmt.Errorf("permission denied") } - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte("binary data")) + body := []byte("binary data") + checksum := fmt.Sprintf("%x", sha256.Sum256(body)) + + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(body) })) defer srv.Close() - err := selfUpdate(srv.URL, "") + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) if err == nil { t.Error("expected rename error") } @@ -143,14 +196,18 @@ func TestSelfUpdateChecksumMatch(t *testing.T) { body := []byte("test binary content") checksum := fmt.Sprintf("%x", sha256.Sum256(body)) - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write(body) })) defer srv.Close() + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + // This will fail at the rename step (writing to os.Executable path), // but the checksum verification should pass - err := selfUpdate(srv.URL, checksum) + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) if err == nil { t.Error("expected error (can't replace running binary in test)") } @@ -160,3 +217,38 @@ func TestSelfUpdateChecksumMatch(t *testing.T) { t.Logf("got expected post-checksum error: %v", err) } } + +func TestSelfUpdateTooLarge(t *testing.T) { + origMax := maxUpdateSize + defer func() { maxUpdateSize = origMax }() + maxUpdateSize = 100 // 100 bytes + + body := make([]byte, 200) // Larger than limit + checksum := fmt.Sprintf("%x", sha256.Sum256(body)) + + srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write(body) + })) + defer srv.Close() + + origGet := httpGet + defer func() { httpGet = origGet }() + httpGet = srv.Client().Get + + err := selfUpdate(rewriteToHTTPS(srv.URL), checksum) + if err == nil { + t.Error("expected error for oversized download") + } + if !strings.Contains(err.Error(), "exceeds max") { + t.Errorf("expected 'exceeds max' in error, got: %v", err) + } +} + +// rewriteToHTTPS converts an httptest TLS server URL to use the https scheme. +// httptest.NewTLSServer returns URLs with https:// already, but this ensures consistency. +func rewriteToHTTPS(rawURL string) string { + if strings.HasPrefix(rawURL, "http://") { + return "https://" + strings.TrimPrefix(rawURL, "http://") + } + return rawURL +} diff --git a/websocket.go b/websocket.go index cabdae2..2bca9f2 100644 --- a/websocket.go +++ b/websocket.go @@ -20,6 +20,8 @@ const ( opClose = 8 opPing = 9 opPong = 10 + + maxFrameSize = 16 << 20 // 16 MB ) // WSConn is a minimal RFC 6455 WebSocket client. @@ -152,6 +154,10 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) { length = binary.BigEndian.Uint64(ext[:]) } + if length > maxFrameSize { + return 0, nil, fmt.Errorf("frame size %d exceeds max %d", length, maxFrameSize) + } + var maskKey [4]byte if masked { if _, err = io.ReadFull(ws.reader, maskKey[:]); err != nil { diff --git a/websocket_test.go b/websocket_test.go index 10d9605..cdd364d 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -307,6 +307,25 @@ func TestReadMessagePongError(t *testing.T) { } } +func TestReadFrameExceedsMaxSize(t *testing.T) { + var buf bytes.Buffer + buf.WriteByte(0x82) // FIN + binary + buf.WriteByte(127) // 64-bit extended length + var extLen [8]byte + binary.BigEndian.PutUint64(extLen[:], uint64(maxFrameSize+1)) + buf.Write(extLen[:]) + // Don't need to write the payload — should reject before reading it + + ws := testWSConn(&nopCloser{readWriter: &buf}) + _, _, err := ws.readFrame() + if err == nil { + t.Error("expected error for frame exceeding max size") + } + if !strings.Contains(err.Error(), "exceeds max") { + t.Errorf("expected 'exceeds max' in error, got: %v", err) + } +} + func TestReadFrameEmptyPayload(t *testing.T) { var buf bytes.Buffer buf.WriteByte(0x82) // FIN + binary