towerops-agent/update_test.go

package main

import (
	"crypto/sha256"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

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("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.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = w.Write([]byte("fake binary"))
	}))
	defer srv.Close()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	err := selfUpdate(rewriteToHTTPS(srv.URL), "0000000000000000000000000000000000000000000000000000000000000000")
	if err == nil {
		t.Error("expected checksum mismatch error")
	}
}

func TestSelfUpdate404(t *testing.T) {
	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
	}))
	defer srv.Close()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	err := selfUpdate(rewriteToHTTPS(srv.URL)+"/missing", "abc123")
	if err == nil {
		t.Error("expected error for 404 response")
	}
	if err != nil && !strings.Contains(err.Error(), "status 404") {
		t.Errorf("expected 'status 404' in error, got: %v", err)
	}
}

func TestSelfUpdateReadBodyError(t *testing.T) {
	// Server sends Content-Length header but closes connection prematurely
	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"))
		// Connection closes without sending the full body
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
	}))
	defer srv.Close()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	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
}

func TestSelfUpdateOsExecutableError(t *testing.T) {
	origExe := osExecutable
	defer func() { osExecutable = origExe }()
	osExecutable = func() (string, error) {
		return "", fmt.Errorf("executable not found")
	}

	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()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
	if err == nil {
		t.Error("expected os.Executable error")
	}
	if !strings.Contains(err.Error(), "get executable path") {
		t.Errorf("expected 'get executable path' in error, got: %v", err)
	}
}

func TestSelfUpdateCreateTempError(t *testing.T) {
	origExe := osExecutable
	origCreate := osCreateTemp
	defer func() {
		osExecutable = origExe
		osCreateTemp = origCreate
	}()
	osExecutable = func() (string, error) { return "/tmp/test-agent", nil }
	osCreateTemp = func(dir, pattern string) (*os.File, error) {
		return nil, fmt.Errorf("disk full")
	}

	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()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
	if err == nil {
		t.Error("expected create temp error")
	}
	if !strings.Contains(err.Error(), "create temp") {
		t.Errorf("expected 'create temp' in error, got: %v", err)
	}
}

func TestSelfUpdateRenameError(t *testing.T) {
	origExe := osExecutable
	origRename := osRename
	defer func() {
		osExecutable = origExe
		osRename = origRename
	}()
	dir := t.TempDir()
	osExecutable = func() (string, error) { return filepath.Join(dir, "test-agent"), nil }
	osRename = func(oldpath, newpath string) error {
		return fmt.Errorf("permission denied")
	}

	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()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
	if err == nil {
		t.Error("expected rename error")
	}
	if !strings.Contains(err.Error(), "rename") {
		t.Errorf("expected 'rename' in error, got: %v", err)
	}
}

func TestSelfUpdateChecksumMatch(t *testing.T) {
	body := []byte("test binary content")
	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()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	// This will fail at the rename step (writing to os.Executable path),
	// but the checksum verification should pass
	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
	if err == nil {
		t.Error("expected error (can't replace running binary in test)")
	}
	// The error should NOT be about checksum
	if err != nil && err.Error() != "" {
		// As long as it's not a checksum error, the checksum verification passed
		t.Logf("got expected post-checksum error: %v", err)
	}
}

func TestSelfUpdateFilePermissions(t *testing.T) {
	origExe := osExecutable
	origRename := osRename
	defer func() {
		osExecutable = origExe
		osRename = origRename
	}()
	dir := t.TempDir()
	osExecutable = func() (string, error) { return filepath.Join(dir, "test-agent"), nil }

	var capturedPath string
	osRename = func(oldpath, newpath string) error {
		capturedPath = oldpath
		return fmt.Errorf("stop here") // stop before re-exec
	}

	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()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	_ = selfUpdate(rewriteToHTTPS(srv.URL), checksum)

	if capturedPath != "" {
		info, err := os.Stat(capturedPath)
		if err == nil {
			if info.Mode().Perm() != 0700 {
				t.Errorf("expected file permissions 0700, got %o", info.Mode().Perm())
			}
		}
	}
}

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()

	origDo := httpDo
	defer func() { httpDo = origDo }()
	httpDo = srv.Client().Do

	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)
	}
}

func TestSanitizeArgs(t *testing.T) {
	tests := []struct {
		name string
		in   []string
		want []string
	}{
		{"flag with separate value", []string{"agent", "--token", "secret"}, []string{"agent", "--token", "***"}},
		{"no sensitive flags", []string{"agent", "--api-url", "wss://x"}, []string{"agent", "--api-url", "wss://x"}},
		{"equals syntax", []string{"agent", "--token=secret"}, []string{"agent", "--token=***"}},
		{"empty args", []string{"agent"}, []string{"agent"}},
		{"trailing flag no value", []string{"agent", "--token"}, []string{"agent", "--token"}},
		{"short flag with value", []string{"agent", "-token", "secret"}, []string{"agent", "-token", "***"}},
		{"short equals syntax", []string{"agent", "-token=secret"}, []string{"agent", "-token=***"}},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := sanitizeArgs(tt.in)
			if len(got) != len(tt.want) {
				t.Fatalf("length: got %d, want %d", len(got), len(tt.want))
			}
			for i := range got {
				if got[i] != tt.want[i] {
					t.Errorf("arg[%d]: got %q, want %q", i, got[i], tt.want[i])
				}
			}
		})
	}

	// Verify original slice is not mutated
	t.Run("does not mutate input", func(t *testing.T) {
		orig := []string{"agent", "--token", "secret"}
		_ = sanitizeArgs(orig)
		if orig[2] != "secret" {
			t.Error("sanitizeArgs mutated the input slice")
		}
	})
}

func TestSelfUpdateInvalidURL(t *testing.T) {
	err := selfUpdate("://\x7f", "abc123")
	if err == nil {
		t.Error("expected error for invalid URL")
	}
	if err != nil && !strings.Contains(err.Error(), "parse url") {
		t.Errorf("expected 'parse url' in error, got: %v", err)
	}
}

func TestSelfUpdateFullHappyPath(t *testing.T) {
	body := []byte("test binary content for full path")
	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()

	origDo := httpDo
	origExe := osExecutable
	origRename := osRename
	origExec := syscallExec
	defer func() {
		httpDo = origDo
		osExecutable = origExe
		osRename = origRename
		syscallExec = origExec
	}()
	httpDo = srv.Client().Do

	dir := t.TempDir()
	exePath := filepath.Join(dir, "test-agent")
	osExecutable = func() (string, error) { return exePath, nil }
	osRename = func(oldpath, newpath string) error {
		return os.Rename(oldpath, newpath) // real rename within temp dir
	}
	syscallExec = func(argv0 string, argv []string, envv []string) error {
		return nil // success — don't actually re-exec
	}

	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
	if err != nil {
		t.Errorf("expected nil error on full happy path, got: %v", err)
	}
}

func TestSelfUpdateWriteError(t *testing.T) {
	body := []byte("binary")
	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()

	origDo := httpDo
	origExe := osExecutable
	origCreate := osCreateTemp
	defer func() {
		httpDo = origDo
		osExecutable = origExe
		osCreateTemp = origCreate
	}()
	httpDo = srv.Client().Do

	dir := t.TempDir()
	osExecutable = func() (string, error) { return filepath.Join(dir, "agent"), nil }

	// Create a temp file that will fail on Write by closing it before selfUpdate tries to write
	osCreateTemp = func(d, pattern string) (*os.File, error) {
		f, err := os.CreateTemp(d, pattern)
		if err != nil {
			return nil, err
		}
		// Close the file so Write will fail
		_ = f.Close()
		return f, nil
	}

	err := selfUpdate(rewriteToHTTPS(srv.URL), checksum)
	if err == nil {
		t.Error("expected write error")
	}
}

func TestSelfUpdateDownloadTimeout(t *testing.T) {
	origDo := httpDo
	origTimeout := selfUpdateTimeout
	defer func() {
		httpDo = origDo
		selfUpdateTimeout = origTimeout
	}()

	httpDo = func(req *http.Request) (*http.Response, error) {
		<-req.Context().Done()
		return nil, req.Context().Err()
	}
	selfUpdateTimeout = 10 * time.Millisecond

	err := selfUpdate("https://example.com/agent", strings.Repeat("0", 64))
	if err == nil {
		t.Fatal("expected timeout error")
	}
	if !strings.Contains(err.Error(), "context deadline exceeded") {
		t.Fatalf("expected context deadline exceeded, 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
}