diff --git a/.gitignore b/.gitignore index 39b99f1..1ca3308 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ towerops-agent *.test cover.out +cover.html # Database files *.db diff --git a/agent.go b/agent.go index e8bf25b..1acc223 100644 --- a/agent.go +++ b/agent.go @@ -17,6 +17,9 @@ import ( "google.golang.org/protobuf/proto" ) +var osExit = os.Exit +var doSelfUpdate = selfUpdate + // channelMsg is the WebSocket channel message format (JSON wrapper around binary protobuf). type channelMsg struct { Topic string `json:"topic"` @@ -267,7 +270,7 @@ func handleMessage( case "restart": slog.Info("restart requested by server, exiting") - os.Exit(0) + osExit(0) case "update": var payload struct { @@ -279,7 +282,7 @@ func handleMessage( return } slog.Info("update requested", "url", payload.URL) - if err := selfUpdate(payload.URL, payload.Checksum); err != nil { + if err := doSelfUpdate(payload.URL, payload.Checksum); err != nil { slog.Error("self-update failed", "error", err) } diff --git a/agent_test.go b/agent_test.go index f6358f7..124125b 100644 --- a/agent_test.go +++ b/agent_test.go @@ -1,8 +1,15 @@ package main import ( + "encoding/base64" "encoding/json" + "fmt" "testing" + "time" + + "github.com/gosnmp/gosnmp" + "github.com/towerops-app/towerops-agent/pb" + "google.golang.org/protobuf/proto" ) func TestChannelMsgSerialization(t *testing.T) { @@ -67,3 +74,345 @@ func searchString(s, substr string) bool { } return false } + +// makeJobPayload creates a base64-encoded protobuf job list payload. +func makeJobPayload(jobs ...*pb.AgentJob) json.RawMessage { + list := &pb.AgentJobList{Jobs: jobs} + bin, _ := proto.Marshal(list) + payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString(bin)}) + return payload +} + +func TestHandleMessage(t *testing.T) { + t.Run("phx_reply", func(t *testing.T) { + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + handleMessage(channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, snmpCh, mtCh, credCh, monCh) + // Just verify it doesn't panic + }) + + t.Run("jobs valid protobuf", func(t *testing.T) { + origDial := snmpDial + defer func() { snmpDial = origDial }() + + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + }, func() {}, 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 := makeJobPayload(&pb.AgentJob{ + JobId: "j1", + JobType: pb.JobType_POLL, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + }) + + handleMessage(channelMsg{Event: "jobs", Payload: payload}, snmpCh, mtCh, credCh, monCh) + // Wait for goroutine to finish + select { + case <-snmpCh: + case <-time.After(2 * time.Second): + t.Error("timed out waiting for snmp result") + } + }) + + t.Run("invalid payload json", func(t *testing.T) { + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + handleMessage(channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, snmpCh, mtCh, credCh, monCh) + // Should log error but not panic + }) + + t.Run("invalid base64", func(t *testing.T) { + 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{"binary": "not-base64!!!"}) + handleMessage(channelMsg{Event: "jobs", Payload: payload}, snmpCh, mtCh, credCh, monCh) + }) + + t.Run("invalid protobuf", func(t *testing.T) { + 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{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})}) + handleMessage(channelMsg{Event: "jobs", Payload: payload}, snmpCh, mtCh, credCh, monCh) + }) + + t.Run("restart", func(t *testing.T) { + origExit := osExit + defer func() { osExit = origExit }() + + var exitCode int + osExit = func(code int) { exitCode = code } + + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + handleMessage(channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, snmpCh, mtCh, credCh, monCh) + + if exitCode != 0 { + t.Errorf("expected exit code 0, got %d", exitCode) + } + }) + + t.Run("update success", func(t *testing.T) { + origUpdate := doSelfUpdate + defer func() { doSelfUpdate = origUpdate }() + + var calledURL string + doSelfUpdate = func(url, checksum string) error { + calledURL = url + 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", "checksum": "abc123"}) + handleMessage(channelMsg{Event: "update", Payload: payload}, snmpCh, mtCh, credCh, monCh) + + if calledURL != "https://example.com/agent" { + t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL) + } + }) + + t.Run("update invalid payload", 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) + // Missing URL field + payload, _ := json.Marshal(map[string]string{"checksum": "abc123"}) + handleMessage(channelMsg{Event: "update", Payload: payload}, snmpCh, mtCh, credCh, monCh) + + if called { + t.Error("selfUpdate should not be called with empty URL") + } + }) + + t.Run("update error", func(t *testing.T) { + origUpdate := doSelfUpdate + defer func() { doSelfUpdate = origUpdate }() + + doSelfUpdate = func(url, checksum string) error { + return fmt.Errorf("download failed") + } + + 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(channelMsg{Event: "update", Payload: payload}, snmpCh, mtCh, credCh, monCh) + // Should log error but not panic + }) + + t.Run("unknown event", func(t *testing.T) { + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + handleMessage(channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, snmpCh, mtCh, credCh, monCh) + // Should just log and not panic + }) + + t.Run("discovery_job event", func(t *testing.T) { + origDial := snmpDial + defer func() { snmpDial = origDial }() + + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + }, func() {}, 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 := makeJobPayload(&pb.AgentJob{ + JobId: "d1", + JobType: pb.JobType_DISCOVER, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }) + handleMessage(channelMsg{Event: "discovery_job", Payload: payload}, snmpCh, mtCh, credCh, monCh) + select { + case <-snmpCh: + case <-time.After(2 * time.Second): + t.Error("timed out waiting for discovery result") + } + }) + + t.Run("backup_job event", func(t *testing.T) { + origDial := mikrotikDial + origSSH := sshBackup + defer func() { mikrotikDial = origDial; sshBackup = origSSH }() + + sshBackup = func(ip string, port uint16, username, password string) (string, error) { + return "/ip address\nadd address=10.0.0.1/24", 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 := makeJobPayload(&pb.AgentJob{ + JobId: "backup:dev1", + JobType: pb.JobType_MIKROTIK, + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, + }) + handleMessage(channelMsg{Event: "backup_job", Payload: payload}, snmpCh, mtCh, credCh, monCh) + select { + case result := <-mtCh: + if result.Error != "" { + t.Errorf("unexpected error: %s", result.Error) + } + case <-time.After(2 * time.Second): + t.Error("timed out waiting for backup result") + } + }) +} + +func TestDispatchJob(t *testing.T) { + t.Run("MIKROTIK", func(t *testing.T) { + origDial := mikrotikDial + defer func() { mikrotikDial = origDial }() + mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { + return nil, fmt.Errorf("not reachable") + } + + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + + dispatchJob(&pb.AgentJob{ + JobId: "mt1", + JobType: pb.JobType_MIKROTIK, + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, + }, snmpCh, mtCh, credCh, monCh) + + select { + case result := <-mtCh: + if result.Error == "" { + t.Error("expected error from unreachable device") + } + case <-time.After(2 * time.Second): + t.Error("timed out") + } + }) + + t.Run("TEST_CREDENTIALS", func(t *testing.T) { + origDial := snmpDial + defer func() { snmpDial = origDial }() + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return nil, nil, fmt.Errorf("refused") + } + + snmpCh := make(chan *pb.SnmpResult, 1) + mtCh := make(chan *pb.MikrotikResult, 1) + credCh := make(chan *pb.CredentialTestResult, 1) + monCh := make(chan *pb.MonitoringCheck, 1) + + dispatchJob(&pb.AgentJob{ + JobId: "tc1", + JobType: pb.JobType_TEST_CREDENTIALS, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, snmpCh, mtCh, credCh, monCh) + + select { + case result := <-credCh: + if result.Success { + t.Error("expected failure") + } + case <-time.After(2 * time.Second): + t.Error("timed out") + } + }) + + t.Run("PING", func(t *testing.T) { + origPing := doPing + defer func() { doPing = origPing }() + doPing = func(ip string, timeoutMs int) (float64, error) { + return 5.5, 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) + + dispatchJob(&pb.AgentJob{ + JobId: "p1", + JobType: pb.JobType_PING, + SnmpDevice: &pb.SnmpDevice{Ip: "127.0.0.1"}, + }, snmpCh, mtCh, credCh, monCh) + + select { + case result := <-monCh: + if result.Status != "success" { + t.Errorf("expected success, got %q", result.Status) + } + case <-time.After(2 * time.Second): + t.Error("timed out") + } + }) + + t.Run("default SNMP", func(t *testing.T) { + origDial := snmpDial + defer func() { snmpDial = origDial }() + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + }, func() {}, 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) + + dispatchJob(&pb.AgentJob{ + JobId: "s1", + JobType: pb.JobType_POLL, + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, snmpCh, mtCh, credCh, monCh) + + select { + case <-snmpCh: + case <-time.After(2 * time.Second): + t.Error("timed out") + } + }) +} diff --git a/log.go b/log.go new file mode 100644 index 0000000..52fcad7 --- /dev/null +++ b/log.go @@ -0,0 +1,118 @@ +package main + +import ( + "context" + "io" + "log/slog" + "sync" +) + +// colorHandler is a slog.Handler that colorizes the level label. +type colorHandler struct { + w io.Writer + level slog.Level + mu sync.Mutex + attrs []slog.Attr + group string +} + +const ( + colorReset = "\033[0m" + colorCyan = "\033[36m" + colorGreen = "\033[32m" + colorYellow = "\033[33m" + colorRed = "\033[31m" +) + +func newColorHandler(w io.Writer, opts *slog.HandlerOptions) *colorHandler { + level := slog.LevelInfo + if opts != nil { + level = opts.Level.Level() + } + return &colorHandler{w: w, level: level} +} + +func (h *colorHandler) Enabled(_ context.Context, level slog.Level) bool { + return level >= h.level +} + +func levelColor(level slog.Level) string { + switch { + case level >= slog.LevelError: + return colorRed + case level >= slog.LevelWarn: + return colorYellow + case level >= slog.LevelInfo: + return colorGreen + default: + return colorCyan + } +} + +func (h *colorHandler) Handle(_ context.Context, r slog.Record) error { + color := levelColor(r.Level) + levelStr := r.Level.String() + + h.mu.Lock() + defer h.mu.Unlock() + + buf := make([]byte, 0, 256) + buf = append(buf, r.Time.Format("2006/01/02 15:04:05")...) + buf = append(buf, ' ') + buf = append(buf, color...) + buf = append(buf, levelStr...) + buf = append(buf, colorReset...) + buf = append(buf, ' ') + buf = append(buf, r.Message...) + + for _, a := range h.attrs { + buf = append(buf, ' ') + buf = appendAttr(buf, h.group, a) + } + + r.Attrs(func(a slog.Attr) bool { + buf = append(buf, ' ') + buf = appendAttr(buf, h.group, a) + return true + }) + + buf = append(buf, '\n') + _, err := h.w.Write(buf) + return err +} + +func appendAttr(buf []byte, group string, a slog.Attr) []byte { + if group != "" { + buf = append(buf, group...) + buf = append(buf, '.') + } + buf = append(buf, a.Key...) + buf = append(buf, '=') + buf = append(buf, a.Value.String()...) + return buf +} + +func (h *colorHandler) WithAttrs(attrs []slog.Attr) slog.Handler { + newAttrs := make([]slog.Attr, len(h.attrs)+len(attrs)) + copy(newAttrs, h.attrs) + copy(newAttrs[len(h.attrs):], attrs) + return &colorHandler{ + w: h.w, + level: h.level, + attrs: newAttrs, + group: h.group, + } +} + +func (h *colorHandler) WithGroup(name string) slog.Handler { + g := name + if h.group != "" { + g = h.group + "." + name + } + return &colorHandler{ + w: h.w, + level: h.level, + attrs: h.attrs, + group: g, + } +} diff --git a/log_test.go b/log_test.go new file mode 100644 index 0000000..f568c4f --- /dev/null +++ b/log_test.go @@ -0,0 +1,159 @@ +package main + +import ( + "bytes" + "context" + "log/slog" + "strings" + "testing" + "time" +) + +func TestLevelColor(t *testing.T) { + tests := []struct { + level slog.Level + want string + }{ + {slog.LevelDebug, colorCyan}, + {slog.LevelInfo, colorGreen}, + {slog.LevelWarn, colorYellow}, + {slog.LevelError, colorRed}, + } + for _, tt := range tests { + got := levelColor(tt.level) + if got != tt.want { + t.Errorf("levelColor(%v) = %q, want %q", tt.level, got, tt.want) + } + } +} + +func TestNewColorHandler(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}) + if h == nil { + t.Fatal("expected non-nil handler") + } +} + +func TestNewColorHandlerNilOpts(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, nil) + if h == nil { + t.Fatal("expected non-nil handler") + } + // Default level should be Info + if h.level != slog.LevelInfo { + t.Errorf("default level: got %v, want %v", h.level, slog.LevelInfo) + } +} + +func TestColorHandlerEnabled(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn}) + + if h.Enabled(context.Background(), slog.LevelDebug) { + t.Error("expected Debug to be disabled with Warn level") + } + if !h.Enabled(context.Background(), slog.LevelWarn) { + t.Error("expected Warn to be enabled") + } + if !h.Enabled(context.Background(), slog.LevelError) { + t.Error("expected Error to be enabled") + } +} + +func TestColorHandlerHandle(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}) + + r := slog.NewRecord(time.Date(2025, 1, 15, 10, 30, 0, 0, time.UTC), slog.LevelInfo, "test message", 0) + r.AddAttrs(slog.String("key", "value")) + + err := h.Handle(context.Background(), r) + if err != nil { + t.Fatal(err) + } + + output := buf.String() + if !strings.Contains(output, "test message") { + t.Errorf("expected 'test message' in output, got: %s", output) + } + if !strings.Contains(output, "key=value") { + t.Errorf("expected 'key=value' in output, got: %s", output) + } + if !strings.Contains(output, colorGreen) { + t.Errorf("expected green color for INFO level in output") + } +} + +func TestColorHandlerWithAttrs(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}) + + h2 := h.WithAttrs([]slog.Attr{slog.String("component", "test")}) + + r := slog.NewRecord(time.Date(2025, 1, 15, 10, 30, 0, 0, time.UTC), slog.LevelInfo, "msg", 0) + err := h2.Handle(context.Background(), r) + if err != nil { + t.Fatal(err) + } + + output := buf.String() + if !strings.Contains(output, "component=test") { + t.Errorf("expected 'component=test' in output, got: %s", output) + } +} + +func TestColorHandlerWithGroup(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}) + + h2 := h.WithGroup("mygroup") + + r := slog.NewRecord(time.Date(2025, 1, 15, 10, 30, 0, 0, time.UTC), slog.LevelInfo, "msg", 0) + r.AddAttrs(slog.String("key", "val")) + err := h2.Handle(context.Background(), r) + if err != nil { + t.Fatal(err) + } + + output := buf.String() + if !strings.Contains(output, "mygroup.key=val") { + t.Errorf("expected 'mygroup.key=val' in output, got: %s", output) + } +} + +func TestColorHandlerWithGroupNested(t *testing.T) { + var buf bytes.Buffer + h := newColorHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}) + + h2 := h.WithGroup("outer").WithGroup("inner") + + r := slog.NewRecord(time.Date(2025, 1, 15, 10, 30, 0, 0, time.UTC), slog.LevelInfo, "msg", 0) + r.AddAttrs(slog.String("k", "v")) + err := h2.Handle(context.Background(), r) + if err != nil { + t.Fatal(err) + } + + output := buf.String() + if !strings.Contains(output, "outer.inner.k=v") { + t.Errorf("expected 'outer.inner.k=v' in output, got: %s", output) + } +} + +func TestAppendAttr(t *testing.T) { + t.Run("no group", func(t *testing.T) { + buf := appendAttr(nil, "", slog.String("key", "value")) + if string(buf) != "key=value" { + t.Errorf("got %q, want %q", string(buf), "key=value") + } + }) + + t.Run("with group", func(t *testing.T) { + buf := appendAttr(nil, "grp", slog.String("key", "value")) + if string(buf) != "grp.key=value" { + t.Errorf("got %q, want %q", string(buf), "grp.key=value") + } + }) +} diff --git a/main.go b/main.go index 1e62ead..fc5a180 100644 --- a/main.go +++ b/main.go @@ -31,7 +31,7 @@ func main() { default: level = slog.LevelInfo } - slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level}))) + slog.SetDefault(slog.New(newColorHandler(os.Stderr, &slog.HandlerOptions{Level: level}))) if *apiURL == "" || *token == "" { fmt.Fprintln(os.Stderr, "error: --api-url and --token are required (or set TOWEROPS_API_URL and TOWEROPS_AGENT_TOKEN)") diff --git a/main_test.go b/main_test.go index af0b492..cb74a79 100644 --- a/main_test.go +++ b/main_test.go @@ -1,6 +1,26 @@ package main -import "testing" +import ( + "os" + "testing" +) + +func TestEnvOrDefault(t *testing.T) { + key := "TOWEROPS_TEST_ENV_OR_DEFAULT" + + // Unset case + os.Unsetenv(key) + if got := envOrDefault(key, "fallback"); got != "fallback" { + t.Errorf("unset: got %q, want %q", got, "fallback") + } + + // Set case + os.Setenv(key, "custom") + defer os.Unsetenv(key) + if got := envOrDefault(key, "fallback"); got != "custom" { + t.Errorf("set: got %q, want %q", got, "custom") + } +} func TestToWebSocketURL(t *testing.T) { tests := []struct { diff --git a/mikrotik.go b/mikrotik.go index 194ae87..13189f5 100644 --- a/mikrotik.go +++ b/mikrotik.go @@ -18,6 +18,8 @@ const ( mikrotikReadTimeout = 30 * time.Second ) +var mikrotikDial = mikrotikConnect + // mikrotikClient is a RouterOS binary API client. type mikrotikClient struct { conn io.ReadWriteCloser @@ -259,7 +261,7 @@ func executeMikrotikJob(job *pb.AgentJob, resultCh chan<- *pb.MikrotikResult) { slog.Debug("executing mikrotik job", "job_id", job.JobId, "device", dev.Ip, "port", dev.Port, "ssl", dev.UseSsl) - client, err := mikrotikConnect(dev.Ip, dev.Port, dev.Username, dev.Password, dev.UseSsl) + client, err := mikrotikDial(dev.Ip, dev.Port, dev.Username, dev.Password, dev.UseSsl) if err != nil { resultCh <- &pb.MikrotikResult{ DeviceId: job.DeviceId, @@ -307,7 +309,7 @@ func executeMikrotikJob(job *pb.AgentJob, resultCh chan<- *pb.MikrotikResult) { func executeMikrotikBackupViaSSH(job *pb.AgentJob, dev *pb.MikrotikDevice, resultCh chan<- *pb.MikrotikResult, timestamp int64) { slog.Debug("executing backup via ssh", "device", job.DeviceId, "ip", dev.Ip, "ssh_port", dev.SshPort) - config, err := executeMikrotikBackup(dev.Ip, uint16(dev.SshPort), dev.Username, dev.Password) + config, err := sshBackup(dev.Ip, uint16(dev.SshPort), dev.Username, dev.Password) if err != nil { resultCh <- &pb.MikrotikResult{ DeviceId: job.DeviceId, diff --git a/mikrotik_test.go b/mikrotik_test.go index e62af6c..375237c 100644 --- a/mikrotik_test.go +++ b/mikrotik_test.go @@ -1,6 +1,12 @@ package main -import "testing" +import ( + "bytes" + "fmt" + "io" + "net" + "testing" +) func TestEncodeLength(t *testing.T) { tests := []struct { @@ -62,3 +68,543 @@ func TestParseMikrotikAttrs(t *testing.T) { }) } } + +// encodeSentence encodes a list of words into RouterOS binary format. +func encodeSentence(words []string) []byte { + var buf []byte + for _, w := range words { + buf = append(buf, encodeLength(len(w))...) + buf = append(buf, w...) + } + buf = append(buf, 0) // empty word terminates sentence + return buf +} + +func TestReadLength(t *testing.T) { + tests := []struct { + name string + data []byte + want int + }{ + {"1-byte (0)", []byte{0x00}, 0}, + {"1-byte (5)", []byte{0x05}, 5}, + {"1-byte (127)", []byte{0x7F}, 127}, + {"2-byte (128)", []byte{0x80, 0x80}, 128}, + {"2-byte (16383)", []byte{0xBF, 0xFF}, 16383}, + {"3-byte (16384)", []byte{0xC0, 0x40, 0x00}, 16384}, + {"3-byte (2097151)", []byte{0xDF, 0xFF, 0xFF}, 2097151}, + {"4-byte (2097152)", []byte{0xE0, 0x20, 0x00, 0x00}, 2097152}, + {"5-byte", []byte{0xF0, 0x10, 0x00, 0x00, 0x00}, 0x10000000}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := &mikrotikClient{conn: &nopCloser{readWriter: bytes.NewBuffer(tt.data)}} + got, err := c.readLength() + if err != nil { + t.Fatal(err) + } + if got != tt.want { + t.Errorf("got %d, want %d", got, tt.want) + } + }) + } +} + +func TestReadWord(t *testing.T) { + t.Run("normal word", func(t *testing.T) { + word := "!done" + var buf bytes.Buffer + buf.Write(encodeLength(len(word))) + buf.WriteString(word) + c := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + got, err := c.readWord() + if err != nil { + t.Fatal(err) + } + if got != word { + t.Errorf("got %q, want %q", got, word) + } + }) + + t.Run("empty word", func(t *testing.T) { + buf := bytes.NewBuffer([]byte{0x00}) + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + got, err := c.readWord() + if err != nil { + t.Fatal(err) + } + if got != "" { + t.Errorf("got %q, want empty", got) + } + }) +} + +func TestReadSentence(t *testing.T) { + var buf bytes.Buffer + buf.Write(encodeSentence([]string{"!re", "=name=eth0", "=type=ether"})) + c := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + words, err := c.readSentence() + if err != nil { + t.Fatal(err) + } + if len(words) != 3 { + t.Fatalf("got %d words, want 3", len(words)) + } + if words[0] != "!re" || words[1] != "=name=eth0" || words[2] != "=type=ether" { + t.Errorf("unexpected words: %v", words) + } +} + +func TestReadResponse(t *testing.T) { + tests := []struct { + name string + sentences [][]string + wantCount int + wantErr string + wantFatalErr bool + }{ + { + name: "done only", + sentences: [][]string{{"!done"}}, + wantCount: 0, + }, + { + name: "done with attrs", + sentences: [][]string{{"!done", "=ret=ok"}}, + wantCount: 1, + }, + { + name: "re + done", + sentences: [][]string{{"!re", "=name=eth0"}, {"!re", "=name=eth1"}, {"!done"}}, + wantCount: 2, + }, + { + name: "trap + done", + sentences: [][]string{{"!trap", "=message=no such command"}, {"!done"}}, + wantErr: "no such command", + wantCount: 0, + }, + { + name: "trap without message + done", + sentences: [][]string{{"!trap"}, {"!done"}}, + wantErr: "unknown error", + wantCount: 0, + }, + { + name: "fatal", + sentences: [][]string{{"!fatal", "=message=connection reset"}}, + wantFatalErr: true, + }, + { + name: "fatal without message", + sentences: [][]string{{"!fatal"}}, + wantFatalErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var buf bytes.Buffer + for _, s := range tt.sentences { + buf.Write(encodeSentence(s)) + } + c := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + resp, err := c.readResponse() + if tt.wantFatalErr { + if err == nil { + t.Error("expected fatal error") + } + return + } + if err != nil { + t.Fatal(err) + } + if len(resp.sentences) != tt.wantCount { + t.Errorf("got %d sentences, want %d", len(resp.sentences), tt.wantCount) + } + if resp.err != tt.wantErr { + t.Errorf("err: got %q, want %q", resp.err, tt.wantErr) + } + }) + } +} + +func TestWriteSentence(t *testing.T) { + var buf bytes.Buffer + c := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + words := []string{"/interface/print", "=detail="} + if err := c.writeSentence(words); err != nil { + t.Fatal(err) + } + + // Read it back + c2 := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + got, err := c2.readSentence() + if err != nil { + t.Fatal(err) + } + if len(got) != len(words) { + t.Fatalf("got %d words, want %d", len(got), len(words)) + } + for i, w := range words { + if got[i] != w { + t.Errorf("word[%d]: got %q, want %q", i, got[i], w) + } + } +} + +func TestExecute(t *testing.T) { + // Use io.Pipe to simulate a full-duplex connection + clientR, serverW := io.Pipe() + serverR, clientW := io.Pipe() + + conn := &readWriteCloser{r: clientR, w: clientW} + c := &mikrotikClient{conn: conn} + + // Server goroutine: read command, write response + go func() { + defer serverW.Close() + sc := &mikrotikClient{conn: &readWriteCloser{r: serverR, w: serverW}} + // Read the command sentence + _, _ = sc.readSentence() + // Write !done response + _ = sc.writeSentence([]string{"!done", "=ret=ok"}) + }() + + resp, err := c.execute("/system/identity/print", nil) + if err != nil { + t.Fatal(err) + } + if resp.err != "" { + t.Errorf("unexpected error: %s", resp.err) + } + if len(resp.sentences) != 1 { + t.Fatalf("got %d sentences, want 1", len(resp.sentences)) + } + if resp.sentences[0].attributes["ret"] != "ok" { + t.Errorf("got ret=%q, want %q", resp.sentences[0].attributes["ret"], "ok") + } +} + +func TestExecuteWithArgs(t *testing.T) { + clientR, serverW := io.Pipe() + serverR, clientW := io.Pipe() + + conn := &readWriteCloser{r: clientR, w: clientW} + c := &mikrotikClient{conn: conn} + + var receivedWords []string + go func() { + defer serverW.Close() + sc := &mikrotikClient{conn: &readWriteCloser{r: serverR, w: serverW}} + receivedWords, _ = sc.readSentence() + _ = sc.writeSentence([]string{"!done"}) + }() + + args := map[string]string{ + "name": "admin", + "?type": "ether", + ".proplist": "name,type", + } + _, err := c.execute("/interface/print", args) + if err != nil { + t.Fatal(err) + } + + // Verify command word + if len(receivedWords) == 0 || receivedWords[0] != "/interface/print" { + t.Errorf("expected command /interface/print, got: %v", receivedWords) + } + + // Verify args formatting: ?-prefix and .-prefix get k=v, others get =k=v + wordSet := make(map[string]bool) + for _, w := range receivedWords[1:] { + wordSet[w] = true + } + if !wordSet["=name=admin"] { + t.Error("expected =name=admin in words") + } + if !wordSet["?type=ether"] { + t.Error("expected ?type=ether in words") + } + if !wordSet[".proplist=name,type"] { + t.Error("expected .proplist=name,type in words") + } +} + +func TestMikrotikClose(t *testing.T) { + clientR, serverW := io.Pipe() + serverR, clientW := io.Pipe() + + conn := &readWriteCloser{r: clientR, w: clientW} + c := &mikrotikClient{conn: conn} + + var receivedWords []string + done := make(chan struct{}) + go func() { + defer close(done) + defer serverW.Close() + sc := &mikrotikClient{conn: &readWriteCloser{r: serverR, w: serverW}} + receivedWords, _ = sc.readSentence() + _ = sc.writeSentence([]string{"!fatal"}) + }() + + _ = c.close() + <-done + + if len(receivedWords) == 0 || receivedWords[0] != "/quit" { + t.Errorf("expected /quit command, got: %v", receivedWords) + } +} + +func TestMikrotikConnect(t *testing.T) { + // Start a test TCP server that speaks mikrotik binary protocol + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + sc := &mikrotikClient{conn: conn} + // Read the /login command + _, _ = sc.readSentence() + // Respond with !done (login success) + _ = sc.writeSentence([]string{"!done"}) + // Read the /quit command on close + _, _ = sc.readSentence() + // Respond with !fatal to close + _ = sc.writeSentence([]string{"!fatal"}) + }() + + _, port, _ := net.SplitHostPort(ln.Addr().String()) + var portNum uint32 + fmt.Sscanf(port, "%d", &portNum) + + client, err := mikrotikConnect("127.0.0.1", portNum, "admin", "pass", false) + if err != nil { + t.Fatal(err) + } + _ = client.close() +} + +func TestMikrotikConnectAuthError(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + sc := &mikrotikClient{conn: conn} + _, _ = sc.readSentence() + // Respond with trap error + done + _ = sc.writeSentence([]string{"!trap", "=message=invalid user"}) + _ = sc.writeSentence([]string{"!done"}) + }() + + _, port, _ := net.SplitHostPort(ln.Addr().String()) + var portNum uint32 + fmt.Sscanf(port, "%d", &portNum) + + _, err = mikrotikConnect("127.0.0.1", portNum, "admin", "wrong", false) + if err == nil { + t.Error("expected auth error") + } +} + +func TestMikrotikConnectFatalError(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + sc := &mikrotikClient{conn: conn} + _, _ = sc.readSentence() + // Respond with fatal error + _ = sc.writeSentence([]string{"!fatal", "=message=connection reset"}) + }() + + _, port, _ := net.SplitHostPort(ln.Addr().String()) + var portNum uint32 + fmt.Sscanf(port, "%d", &portNum) + + _, err = mikrotikConnect("127.0.0.1", portNum, "admin", "pass", false) + if err == nil { + t.Error("expected fatal error") + } +} + +func TestMikrotikConnectRefused(t *testing.T) { + // Connect to a port with nothing listening + _, err := mikrotikConnect("127.0.0.1", 1, "admin", "pass", false) + if err == nil { + t.Error("expected connection refused") + } +} + +func TestMikrotikConnectSSL(t *testing.T) { + // SSL connect to a port with nothing listening — tests the TLS dialer path + _, err := mikrotikConnect("127.0.0.1", 1, "admin", "pass", true) + if err == nil { + t.Error("expected connection error with SSL") + } +} + +func TestReadResponseEmptySentence(t *testing.T) { + // An empty sentence (just the terminator byte) should be skipped + var buf bytes.Buffer + buf.WriteByte(0) // empty sentence + buf.Write(encodeSentence([]string{"!done"})) + c := &mikrotikClient{conn: &nopCloser{readWriter: &buf}} + resp, err := c.readResponse() + if err != nil { + t.Fatal(err) + } + if len(resp.sentences) != 0 { + t.Errorf("expected 0 sentences, got %d", len(resp.sentences)) + } +} + +func TestReadSentenceWithNetConn(t *testing.T) { + // Test readSentence with a real net.Conn to trigger SetReadDeadline path + server, client := net.Pipe() + defer server.Close() + defer client.Close() + + go func() { + server.Write(encodeSentence([]string{"!done"})) + }() + + c := &mikrotikClient{conn: client} + words, err := c.readSentence() + if err != nil { + t.Fatal(err) + } + if len(words) != 1 || words[0] != "!done" { + t.Errorf("unexpected words: %v", words) + } +} + +func TestExecuteWriteError(t *testing.T) { + c := &mikrotikClient{conn: &failWriter{}} + _, err := c.execute("/test", nil) + if err == nil { + t.Error("expected write error") + } +} + +func TestReadWordError(t *testing.T) { + // Empty buffer causes read error + buf := bytes.NewBuffer([]byte{0x05}) // length 5, but no data following + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readWord() + if err == nil { + t.Error("expected read error for truncated word") + } +} + +func TestReadLengthError(t *testing.T) { + // Empty buffer causes EOF + buf := bytes.NewBuffer(nil) + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readLength() + if err == nil { + t.Error("expected read error for empty buffer") + } +} + +func TestReadLength2ByteError(t *testing.T) { + // 2-byte length with truncated second byte + buf := bytes.NewBuffer([]byte{0x80}) // needs 1 more byte + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readLength() + if err == nil { + t.Error("expected error for truncated 2-byte length") + } +} + +func TestReadLength3ByteError(t *testing.T) { + buf := bytes.NewBuffer([]byte{0xC0}) // needs 2 more bytes + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readLength() + if err == nil { + t.Error("expected error for truncated 3-byte length") + } +} + +func TestReadLength4ByteError(t *testing.T) { + buf := bytes.NewBuffer([]byte{0xE0}) // needs 3 more bytes + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readLength() + if err == nil { + t.Error("expected error for truncated 4-byte length") + } +} + +func TestReadLength5ByteError(t *testing.T) { + buf := bytes.NewBuffer([]byte{0xF0}) // needs 4 more bytes + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readLength() + if err == nil { + t.Error("expected error for truncated 5-byte length") + } +} + +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 + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readSentence() + if err == nil { + t.Error("expected error for truncated sentence") + } +} + +func TestReadResponseReadError(t *testing.T) { + // Empty buffer causes immediate error + buf := bytes.NewBuffer(nil) + c := &mikrotikClient{conn: &nopCloser{readWriter: buf}} + _, err := c.readResponse() + if err == nil { + t.Error("expected error for empty buffer") + } +} + +// failWriter always returns an error on Write. +type failWriter struct{} + +func (f *failWriter) Read(p []byte) (int, error) { return 0, io.EOF } +func (f *failWriter) Write(p []byte) (int, error) { return 0, fmt.Errorf("write failed") } +func (f *failWriter) Close() error { return nil } + +// readWriteCloser combines separate reader and writer into io.ReadWriteCloser. +type readWriteCloser struct { + r io.Reader + w io.Writer +} + +func (rwc *readWriteCloser) Read(p []byte) (int, error) { return rwc.r.Read(p) } +func (rwc *readWriteCloser) Write(p []byte) (int, error) { return rwc.w.Write(p) } +func (rwc *readWriteCloser) Close() error { + if c, ok := rwc.w.(io.Closer); ok { + return c.Close() + } + return nil +} diff --git a/ping_test.go b/ping_test.go index 590e3cc..ce4818b 100644 --- a/ping_test.go +++ b/ping_test.go @@ -1,6 +1,42 @@ package main -import "testing" +import ( + "runtime" + "testing" +) + +func TestPingDeviceLocalhost(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping ping test on windows") + } + ms, err := pingDevice("127.0.0.1", 5000) + if err != nil { + t.Fatalf("ping localhost failed: %v", err) + } + if ms <= 0 { + t.Errorf("expected positive response time, got %v", ms) + } +} + +func TestPingDeviceInvalidIP(t *testing.T) { + _, err := pingDevice("not-an-ip", 5000) + if err == nil { + t.Error("expected error for invalid IP") + } +} + +func TestPingDeviceIPv6(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping ping test on windows") + } + ms, err := pingDevice("::1", 5000) + if err != nil { + t.Skipf("IPv6 not available: %v", err) + } + if ms <= 0 { + t.Errorf("expected positive response time, got %v", ms) + } +} func TestParsePingTime(t *testing.T) { tests := []struct { @@ -34,6 +70,11 @@ func TestParsePingTime(t *testing.T) { output: "", wantErr: true, }, + { + name: "time= without ms suffix", + output: "64 bytes from 10.0.0.1: icmp_seq=1 ttl=64 time=1.234\n", + want: 1.234, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/snmp_test.go b/snmp_test.go index c931132..e1eb9fd 100644 --- a/snmp_test.go +++ b/snmp_test.go @@ -1,9 +1,11 @@ package main import ( + "fmt" "testing" "github.com/gosnmp/gosnmp" + "github.com/towerops-app/towerops-agent/pb" ) func TestSnmpValueToString(t *testing.T) { @@ -72,6 +74,26 @@ func TestSnmpValueToString(t *testing.T) { pdu: gosnmp.SnmpPDU{Type: gosnmp.OctetString, Value: []byte{0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x80, 0xFE}}, want: "48:65:6c:6c:6f:80:fe", }, + { + name: "opaque", + pdu: gosnmp.SnmpPDU{Type: gosnmp.Opaque, Value: []byte{0xDE, 0xAD}}, + want: "de:ad", + }, + { + name: "end of mib view", + pdu: gosnmp.SnmpPDU{Type: gosnmp.EndOfMibView, Value: nil}, + want: "null", + }, + { + name: "no such instance", + pdu: gosnmp.SnmpPDU{Type: gosnmp.NoSuchInstance, Value: nil}, + want: "null", + }, + { + name: "unknown type", + pdu: gosnmp.SnmpPDU{Type: gosnmp.Asn1BER(0xFF), Value: "something"}, + want: "something", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -90,7 +112,10 @@ func TestMapAuthProtocol(t *testing.T) { }{ {"MD5", gosnmp.MD5}, {"SHA", gosnmp.SHA}, + {"SHA-1", gosnmp.SHA}, + {"SHA-224", gosnmp.SHA224}, {"SHA-256", gosnmp.SHA256}, + {"SHA-384", gosnmp.SHA384}, {"SHA-512", gosnmp.SHA512}, {"unknown", gosnmp.SHA}, } @@ -109,7 +134,11 @@ func TestMapPrivProtocol(t *testing.T) { }{ {"DES", gosnmp.DES}, {"AES", gosnmp.AES}, + {"AES-128", gosnmp.AES}, + {"AES-192", gosnmp.AES192}, {"AES-256", gosnmp.AES256}, + {"AES-192-C", gosnmp.AES192C}, + {"AES-256-C", gosnmp.AES256C}, {"unknown", gosnmp.AES}, } for _, tt := range tests { @@ -137,3 +166,475 @@ func TestFormatHex(t *testing.T) { } } } + +func TestNewSnmpConn(t *testing.T) { + tests := []struct { + name string + dev *pb.SnmpDevice + }{ + { + name: "v1", + dev: &pb.SnmpDevice{Ip: "127.0.0.1", Port: 0, Version: "v1", Community: "public"}, + }, + { + name: "v2c default", + dev: &pb.SnmpDevice{Ip: "127.0.0.1", Port: 0, Version: "2c", Community: "public"}, + }, + { + name: "v2c empty version", + dev: &pb.SnmpDevice{Ip: "127.0.0.1", Port: 0, Community: "public"}, + }, + { + name: "v3 noAuthNoPriv", + dev: &pb.SnmpDevice{Ip: "127.0.0.1", Port: 0, Version: "v3", V3Username: "user", V3SecurityLevel: "noAuthNoPriv"}, + }, + { + name: "v3 authNoPriv", + dev: &pb.SnmpDevice{ + Ip: "127.0.0.1", Port: 0, Version: "v3", + V3Username: "user", V3SecurityLevel: "authNoPriv", + V3AuthProtocol: "SHA-256", V3AuthPassword: "pass1234", + }, + }, + { + name: "v3 authPriv", + dev: &pb.SnmpDevice{ + Ip: "127.0.0.1", Port: 0, Version: "v3", + V3Username: "user", V3SecurityLevel: "authPriv", + V3AuthProtocol: "SHA-256", V3AuthPassword: "pass1234", + V3PrivProtocol: "AES-256", V3PrivPassword: "priv1234", + }, + }, + { + name: "tcp transport", + dev: &pb.SnmpDevice{Ip: "127.0.0.1", Port: 0, Version: "2c", Community: "public", Transport: "tcp"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + conn, err := newSnmpConn(tt.dev) + if err != nil { + return + } + defer func() { _ = conn.Conn.Close() }() + + // Verify version was set correctly + switch tt.dev.Version { + case "1", "v1": + if conn.Version != gosnmp.Version1 { + t.Errorf("expected Version1, got %v", conn.Version) + } + case "3", "v3": + if conn.Version != gosnmp.Version3 { + t.Errorf("expected Version3, got %v", conn.Version) + } + default: + if conn.Version != gosnmp.Version2c { + t.Errorf("expected Version2c, got %v", conn.Version) + } + } + }) + } +} + +func TestNewSnmpConnTCPError(t *testing.T) { + // TCP transport on port 1 should fail to connect + _, err := newSnmpConn(&pb.SnmpDevice{Ip: "127.0.0.1", Port: 1, Version: "2c", Community: "public", Transport: "tcp"}) + if err == nil { + t.Error("expected connection error on TCP port 1") + } +} + +func TestSnmpDialDefault(t *testing.T) { + // Test the default snmpDial function variable (wraps newSnmpConn) + origDial := snmpDial + defer func() { snmpDial = origDial }() + + // Reset to default behavior + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + conn, err := newSnmpConn(dev) + if err != nil { + return nil, nil, err + } + return conn, func() { _ = conn.Conn.Close() }, nil + } + + q, closeFn, err := snmpDial(&pb.SnmpDevice{Ip: "127.0.0.1", Port: 16100, Version: "2c", Community: "public"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + defer closeFn() + if q == nil { + t.Error("expected non-nil querier") + } +} + +// mockSnmpQuerier implements snmpQuerier for testing. +type mockSnmpQuerier struct { + getFunc func(oids []string) (*gosnmp.SnmpPacket, error) + walkFunc func(rootOid string) ([]gosnmp.SnmpPDU, error) + closeCalled bool +} + +func (m *mockSnmpQuerier) Get(oids []string) (*gosnmp.SnmpPacket, error) { + return m.getFunc(oids) +} + +func (m *mockSnmpQuerier) BulkWalkAll(rootOid string) ([]gosnmp.SnmpPDU, error) { + return m.walkFunc(rootOid) +} + +func TestExecuteSnmpJob(t *testing.T) { + t.Run("nil device", func(t *testing.T) { + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{JobId: "1"}, ch) + if len(ch) != 0 { + t.Error("expected no result for nil device") + } + }) + + t.Run("dial error", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return nil, nil, fmt.Errorf("connection refused") + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + }, ch) + if len(ch) != 0 { + t.Error("expected no result on dial error") + } + }) + + t.Run("GET success", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{ + Variables: []gosnmp.SnmpPDU{ + {Name: ".1.3.6.1.2.1.1.1.0", Type: gosnmp.OctetString, Value: []byte("Linux")}, + }, + }, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() { mock.closeCalled = true }, nil + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + DeviceId: "dev-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_GET, Oids: []string{".1.3.6.1.2.1.1.1.0"}}, + }, + }, ch) + + if len(ch) != 1 { + t.Fatal("expected one result") + } + result := <-ch + if result.OidValues[".1.3.6.1.2.1.1.1.0"] != "Linux" { + t.Errorf("got %q, want Linux", result.OidValues[".1.3.6.1.2.1.1.1.0"]) + } + if !mock.closeCalled { + t.Error("expected close to be called") + } + }) + + t.Run("WALK success", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + walkFunc: func(rootOid string) ([]gosnmp.SnmpPDU, error) { + return []gosnmp.SnmpPDU{ + {Name: ".1.3.6.1.2.1.2.2.1.1.1", Type: gosnmp.Integer, Value: 1}, + {Name: ".1.3.6.1.2.1.2.2.1.1.2", Type: gosnmp.Integer, Value: 2}, + }, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_WALK, Oids: []string{".1.3.6.1.2.1.2.2.1.1"}}, + }, + }, ch) + + result := <-ch + if len(result.OidValues) != 2 { + t.Errorf("got %d oid values, want 2", len(result.OidValues)) + } + }) + + t.Run("GET error continues", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return nil, fmt.Errorf("timeout") + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_GET, Oids: []string{".1.3.6.1.2.1.1.1.0"}}, + }, + }, ch) + + result := <-ch + if len(result.OidValues) != 0 { + t.Errorf("got %d oid values, want 0 on error", len(result.OidValues)) + } + }) + + t.Run("WALK error continues", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + walkFunc: func(rootOid string) ([]gosnmp.SnmpPDU, error) { + return nil, fmt.Errorf("timeout") + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_WALK, Oids: []string{".1.3.6.1.2.1.2"}}, + }, + }, ch) + + result := <-ch + if len(result.OidValues) != 0 { + t.Errorf("got %d oid values, want 0 on error", len(result.OidValues)) + } + }) + + t.Run("NoSuchObject skipped", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{ + Variables: []gosnmp.SnmpPDU{ + {Name: ".1.3.6.1.2.1.1.1.0", Type: gosnmp.NoSuchObject, Value: nil}, + }, + }, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_GET, Oids: []string{".1.3.6.1.2.1.1.1.0"}}, + }, + }, ch) + + result := <-ch + if len(result.OidValues) != 0 { + t.Errorf("NoSuchObject should be skipped, got %d oid values", len(result.OidValues)) + } + }) + + t.Run("WALK NoSuchObject skipped", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + walkFunc: func(rootOid string) ([]gosnmp.SnmpPDU, error) { + return []gosnmp.SnmpPDU{ + {Name: ".1.3.6.1.2.1.2.2.1.1.1", Type: gosnmp.Integer, Value: 1}, + {Name: ".1.3.6.1.2.1.2.2.1.1.2", Type: gosnmp.NoSuchInstance, Value: nil}, + {Name: ".1.3.6.1.2.1.2.2.1.1.3", Type: gosnmp.EndOfMibView, Value: nil}, + }, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.SnmpResult, 1) + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + Queries: []*pb.SnmpQuery{ + {QueryType: pb.QueryType_WALK, Oids: []string{".1.3.6.1.2.1.2.2.1.1"}}, + }, + }, ch) + + result := <-ch + if len(result.OidValues) != 1 { + t.Errorf("expected 1 value (others skipped), got %d", len(result.OidValues)) + } + }) + + t.Run("channel full drops", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{}, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.SnmpResult) // unbuffered, no reader — will be full + executeSnmpJob(&pb.AgentJob{ + JobId: "1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, ch) + // Should not block — the result is dropped + }) +} + +func TestExecuteCredentialTest(t *testing.T) { + t.Run("nil device", func(t *testing.T) { + ch := make(chan *pb.CredentialTestResult, 1) + executeCredentialTest(&pb.AgentJob{JobId: "1"}, ch) + if len(ch) != 0 { + t.Error("expected no result for nil device") + } + }) + + t.Run("dial error", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return nil, nil, fmt.Errorf("connection refused") + } + + ch := make(chan *pb.CredentialTestResult, 1) + executeCredentialTest(&pb.AgentJob{ + JobId: "test-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161}, + }, ch) + + result := <-ch + if result.Success { + t.Error("expected failure") + } + if result.ErrorMessage == "" { + t.Error("expected error message") + } + }) + + t.Run("get error", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return nil, fmt.Errorf("timeout") + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.CredentialTestResult, 1) + executeCredentialTest(&pb.AgentJob{ + JobId: "test-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, ch) + + result := <-ch + if result.Success { + t.Error("expected failure on get error") + } + }) + + t.Run("success", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{ + Variables: []gosnmp.SnmpPDU{ + {Name: ".1.3.6.1.2.1.1.1.0", Type: gosnmp.OctetString, Value: []byte("RouterOS 7.1")}, + }, + }, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.CredentialTestResult, 1) + executeCredentialTest(&pb.AgentJob{ + JobId: "test-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, ch) + + result := <-ch + if !result.Success { + t.Error("expected success") + } + if result.SystemDescription != "RouterOS 7.1" { + t.Errorf("sysDescr: got %q, want %q", result.SystemDescription, "RouterOS 7.1") + } + }) + + t.Run("success no variables", func(t *testing.T) { + orig := snmpDial + defer func() { snmpDial = orig }() + + mock := &mockSnmpQuerier{ + getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) { + return &gosnmp.SnmpPacket{Variables: nil}, nil + }, + } + snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) { + return mock, func() {}, nil + } + + ch := make(chan *pb.CredentialTestResult, 1) + executeCredentialTest(&pb.AgentJob{ + JobId: "test-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, ch) + + result := <-ch + if !result.Success { + t.Error("expected success even with no variables") + } + if result.SystemDescription != "" { + t.Errorf("expected empty sysDescr, got %q", result.SystemDescription) + } + }) +} diff --git a/ssh.go b/ssh.go index b78482a..ebb2428 100644 --- a/ssh.go +++ b/ssh.go @@ -9,6 +9,9 @@ import ( "golang.org/x/crypto/ssh" ) +var sshBackup = executeMikrotikBackup +var doPing = pingDevice + // executeMikrotikBackup connects via SSH and runs /export compact. func executeMikrotikBackup(ip string, port uint16, username, password string) (string, error) { config := &ssh.ClientConfig{ @@ -52,7 +55,7 @@ func executePingJob(job *pb.AgentJob, resultCh chan<- *pb.MonitoringCheck) { } timestamp := time.Now().Unix() - responseTime, err := pingDevice(dev.Ip, 5000) + responseTime, err := doPing(dev.Ip, 5000) if err != nil { slog.Warn("device down", "device", job.DeviceId, "error", err) diff --git a/ssh_test.go b/ssh_test.go new file mode 100644 index 0000000..2a253cd --- /dev/null +++ b/ssh_test.go @@ -0,0 +1,528 @@ +package main + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "fmt" + "net" + "strings" + "testing" + "time" + + "github.com/towerops-app/towerops-agent/pb" + "golang.org/x/crypto/ssh" +) + +func TestExecutePingJob(t *testing.T) { + t.Run("nil device", func(t *testing.T) { + ch := make(chan *pb.MonitoringCheck, 1) + executePingJob(&pb.AgentJob{JobId: "p1"}, ch) + if len(ch) != 0 { + t.Error("expected no result for nil device") + } + }) + + t.Run("success", func(t *testing.T) { + origPing := doPing + defer func() { doPing = origPing }() + doPing = func(ip string, timeoutMs int) (float64, error) { + return 3.14, nil + } + + ch := make(chan *pb.MonitoringCheck, 1) + executePingJob(&pb.AgentJob{ + JobId: "p1", + DeviceId: "dev-1", + SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"}, + }, ch) + + select { + case result := <-ch: + if result.Status != "success" { + t.Errorf("status: got %q, want %q", result.Status, "success") + } + if result.ResponseTimeMs != 3.14 { + t.Errorf("response time: got %v, want 3.14", result.ResponseTimeMs) + } + if result.DeviceId != "dev-1" { + t.Errorf("device id: got %q, want %q", result.DeviceId, "dev-1") + } + case <-time.After(time.Second): + t.Error("timed out") + } + }) + + t.Run("failure", func(t *testing.T) { + origPing := doPing + defer func() { doPing = origPing }() + doPing = func(ip string, timeoutMs int) (float64, error) { + return 0, fmt.Errorf("request timeout") + } + + ch := make(chan *pb.MonitoringCheck, 1) + executePingJob(&pb.AgentJob{ + JobId: "p2", + DeviceId: "dev-2", + SnmpDevice: &pb.SnmpDevice{Ip: "192.168.1.1"}, + }, ch) + + select { + case result := <-ch: + if result.Status != "failure" { + t.Errorf("status: got %q, want %q", result.Status, "failure") + } + case <-time.After(time.Second): + t.Error("timed out") + } + }) +} + +func TestExecuteMikrotikJob(t *testing.T) { + t.Run("nil device", func(t *testing.T) { + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikJob(&pb.AgentJob{JobId: "m1"}, ch) + if len(ch) != 0 { + t.Error("expected no result for nil device") + } + }) + + t.Run("dial error", func(t *testing.T) { + origDial := mikrotikDial + defer func() { mikrotikDial = origDial }() + mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { + return nil, fmt.Errorf("connection refused") + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikJob(&pb.AgentJob{ + JobId: "m1", + DeviceId: "dev-1", + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, + }, ch) + + result := <-ch + if result.Error == "" { + t.Error("expected error") + } + }) + + t.Run("success", func(t *testing.T) { + origDial := mikrotikDial + defer func() { mikrotikDial = origDial }() + + mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { + return newMockMikrotikClient([]mockMikrotikResponse{ + {resp: &mikrotikResponse{sentences: []mikrotikSentence{{attributes: map[string]string{"name": "ether1"}}}}}, + {resp: &mikrotikResponse{}}, // close /quit + }), nil + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikJob(&pb.AgentJob{ + JobId: "m1", + DeviceId: "dev-1", + MikrotikDevice: &pb.MikrotikDevice{ + Ip: "10.0.0.1", Port: 8728, Username: "admin", Password: "pass", + }, + MikrotikCommands: []*pb.MikrotikCommand{ + {Command: "/interface/print"}, + }, + }, ch) + + result := <-ch + if result.Error != "" { + t.Errorf("unexpected error: %s", result.Error) + } + if len(result.Sentences) != 1 { + t.Fatalf("got %d sentences, want 1", len(result.Sentences)) + } + if result.Sentences[0].Attributes["name"] != "ether1" { + t.Errorf("expected name=ether1, got %v", result.Sentences[0].Attributes) + } + }) + + t.Run("command error", func(t *testing.T) { + origDial := mikrotikDial + defer func() { mikrotikDial = origDial }() + + mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { + return newMockMikrotikClient([]mockMikrotikResponse{ + {err: fmt.Errorf("fatal: connection lost")}, + {resp: &mikrotikResponse{}}, // close + }), nil + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikJob(&pb.AgentJob{ + JobId: "m1", + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, + MikrotikCommands: []*pb.MikrotikCommand{ + {Command: "/system/reboot"}, + }, + }, ch) + + result := <-ch + if result.Error == "" { + t.Error("expected error from failed command") + } + }) + + t.Run("response error", func(t *testing.T) { + origDial := mikrotikDial + defer func() { mikrotikDial = origDial }() + + mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) { + return newMockMikrotikClient([]mockMikrotikResponse{ + {resp: &mikrotikResponse{err: "no such command"}}, + {resp: &mikrotikResponse{}}, // close + }), nil + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikJob(&pb.AgentJob{ + JobId: "m1", + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728}, + MikrotikCommands: []*pb.MikrotikCommand{ + {Command: "/bad/command"}, + }, + }, ch) + + result := <-ch + if result.Error == "" { + t.Error("expected error from response error") + } + }) + + t.Run("backup routing via SSH", func(t *testing.T) { + origSSH := sshBackup + defer func() { sshBackup = origSSH }() + + sshBackup = func(ip string, port uint16, username, password string) (string, error) { + return "/ip address\nadd address=10.0.0.1/24", nil + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikJob(&pb.AgentJob{ + JobId: "backup:dev1", + DeviceId: "dev-1", + MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, + }, ch) + + result := <-ch + if result.Error != "" { + t.Errorf("unexpected error: %s", result.Error) + } + if len(result.Sentences) != 1 { + t.Fatalf("got %d sentences, want 1", len(result.Sentences)) + } + if result.Sentences[0].Attributes["config"] == "" { + t.Error("expected config in attributes") + } + }) +} + +func TestExecuteMikrotikBackupViaSSH(t *testing.T) { + t.Run("success", func(t *testing.T) { + origSSH := sshBackup + defer func() { sshBackup = origSSH }() + + sshBackup = func(ip string, port uint16, username, password string) (string, error) { + return "# test config", nil + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikBackupViaSSH( + &pb.AgentJob{JobId: "backup:1", DeviceId: "d1"}, + &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, + ch, 1000, + ) + + result := <-ch + if result.Error != "" { + t.Errorf("unexpected error: %s", result.Error) + } + if len(result.Sentences) != 1 || result.Sentences[0].Attributes["config"] != "# test config" { + t.Error("expected config sentence") + } + }) + + t.Run("error", func(t *testing.T) { + origSSH := sshBackup + defer func() { sshBackup = origSSH }() + + sshBackup = func(ip string, port uint16, username, password string) (string, error) { + return "", fmt.Errorf("ssh connection refused") + } + + ch := make(chan *pb.MikrotikResult, 1) + executeMikrotikBackupViaSSH( + &pb.AgentJob{JobId: "backup:2", DeviceId: "d2"}, + &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"}, + ch, 1000, + ) + + result := <-ch + if result.Error == "" { + t.Error("expected SSH error") + } + }) +} + +func TestExecuteMikrotikBackupDialError(t *testing.T) { + _, err := executeMikrotikBackup("127.0.0.1", 1, "admin", "pass") + if err == nil { + t.Error("expected SSH dial error") + } +} + +func TestExecuteMikrotikBackupSuccess(t *testing.T) { + addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) { + ch.Write([]byte("# RouterOS config\n/ip address\nadd address=10.0.0.1/24\n")) + ch.CloseWrite() + // Send exit-status 0 + ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{0})) + ch.Close() + }) + defer cleanup() + + _, port, _ := net.SplitHostPort(addr) + var portNum uint16 + fmt.Sscanf(port, "%d", &portNum) + + config, err := executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass") + if err != nil { + t.Fatal(err) + } + if config == "" { + t.Error("expected non-empty config") + } +} + +func TestExecuteMikrotikBackupCommandError(t *testing.T) { + addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) { + // Send exit-status 1 with no output (simulates command failure) + ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{1})) + ch.Close() + }) + defer cleanup() + + _, port, _ := net.SplitHostPort(addr) + var portNum uint16 + fmt.Sscanf(port, "%d", &portNum) + + _, err := executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass") + if err == nil { + t.Error("expected error from failed command") + } +} + +func TestExecuteMikrotikBackupWithOutput(t *testing.T) { + // MikroTik SSH returns output even with non-zero exit code + addr, cleanup := startTestSSHServer(t, func(ch ssh.Channel) { + ch.Write([]byte("# partial config\n")) + ch.CloseWrite() + ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{1})) + ch.Close() + }) + defer cleanup() + + _, port, _ := net.SplitHostPort(addr) + var portNum uint16 + fmt.Sscanf(port, "%d", &portNum) + + config, err := executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass") + if err != nil { + t.Fatalf("expected success when output present despite exit code, got: %v", err) + } + if config == "" { + t.Error("expected non-empty config") + } +} + +func TestExecuteMikrotikBackupSessionError(t *testing.T) { + // SSH server that accepts connection but rejects all channel requests + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatal(err) + } + signer, err := ssh.NewSignerFromKey(key) + if err != nil { + t.Fatal(err) + } + + config := &ssh.ServerConfig{ + PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) { + return nil, nil + }, + } + config.AddHostKey(signer) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + + sconn, chans, reqs, err := ssh.NewServerConn(conn, config) + if err != nil { + return + } + defer sconn.Close() + go ssh.DiscardRequests(reqs) + + // Reject all channel requests to trigger NewSession error + for newChannel := range chans { + newChannel.Reject(ssh.Prohibited, "no sessions allowed") + } + }() + + _, port, _ := net.SplitHostPort(ln.Addr().String()) + var portNum uint16 + fmt.Sscanf(port, "%d", &portNum) + + _, err = executeMikrotikBackup("127.0.0.1", portNum, "admin", "pass") + if err == nil { + t.Error("expected session error") + } + if !strings.Contains(err.Error(), "ssh session") { + t.Errorf("expected 'ssh session' in error, got: %v", err) + } +} + +// startTestSSHServer starts a minimal SSH server for testing and returns its address and cleanup function. +func startTestSSHServer(t *testing.T, handler func(ch ssh.Channel)) (string, func()) { + t.Helper() + + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatal(err) + } + signer, err := ssh.NewSignerFromKey(key) + if err != nil { + t.Fatal(err) + } + + config := &ssh.ServerConfig{ + PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) { + return nil, nil // Accept any password + }, + } + config.AddHostKey(signer) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + + sconn, chans, reqs, err := ssh.NewServerConn(conn, config) + if err != nil { + return + } + defer sconn.Close() + go ssh.DiscardRequests(reqs) + + for newChannel := range chans { + if newChannel.ChannelType() != "session" { + newChannel.Reject(ssh.UnknownChannelType, "unknown channel type") + continue + } + ch, requests, err := newChannel.Accept() + if err != nil { + continue + } + go func() { + for req := range requests { + if req.Type == "exec" { + req.Reply(true, nil) + handler(ch) + return + } + req.Reply(false, nil) + } + }() + } + }() + + return ln.Addr().String(), func() { ln.Close() } +} + +// mockMikrotikResponse pairs a response with an optional error for mock execute calls. +type mockMikrotikResponse struct { + resp *mikrotikResponse + err error +} + +// newMockMikrotikClient creates a mikrotikClient backed by a mock that returns +// canned responses from the provided list, in order. +func newMockMikrotikClient(responses []mockMikrotikResponse) *mikrotikClient { + return &mikrotikClient{conn: &mockMikrotikConn{responses: responses}} +} + +// mockMikrotikConn is a fake io.ReadWriteCloser that the mikrotikClient can use. +// It intercepts execute() calls by providing pre-encoded binary responses. +// Since mikrotikClient.execute calls writeSentence then readResponse, we need +// a conn that absorbs writes and returns pre-built binary sentences on read. +type mockMikrotikConn struct { + responses []mockMikrotikResponse + callIdx int + readBuf []byte +} + +func (m *mockMikrotikConn) Write(p []byte) (int, error) { + // Absorb writes (the command sentence). When a full sentence is written, + // prepare the response for the next read. + // We detect sentence end by looking for the 0x00 terminator. + for _, b := range p { + if b == 0x00 { + // A sentence was completed. Prepare the response. + if m.callIdx < len(m.responses) { + r := m.responses[m.callIdx] + m.callIdx++ + if r.err != nil { + // Encode a !fatal response + m.readBuf = append(m.readBuf, encodeSentence([]string{"!fatal", "=message=" + r.err.Error()})...) + } else { + // Encode sentences + for _, s := range r.resp.sentences { + words := []string{"!re"} + for k, v := range s.attributes { + words = append(words, "="+k+"="+v) + } + m.readBuf = append(m.readBuf, encodeSentence(words)...) + } + if r.resp.err != "" { + m.readBuf = append(m.readBuf, encodeSentence([]string{"!trap", "=message=" + r.resp.err})...) + } + m.readBuf = append(m.readBuf, encodeSentence([]string{"!done"})...) + } + } + } + } + return len(p), nil +} + +func (m *mockMikrotikConn) Read(p []byte) (int, error) { + if len(m.readBuf) == 0 { + return 0, fmt.Errorf("no data") + } + n := copy(p, m.readBuf) + m.readBuf = m.readBuf[n:] + return n, nil +} + +func (m *mockMikrotikConn) Close() error { return nil } diff --git a/update.go b/update.go index 6d24ae6..cdca9af 100644 --- a/update.go +++ b/update.go @@ -10,6 +10,10 @@ import ( "syscall" ) +var osExecutable = os.Executable +var osWriteFile = os.WriteFile +var osRename = os.Rename + // selfUpdate downloads a new binary, verifies its checksum, replaces the current binary, and re-execs. func selfUpdate(downloadURL, expectedChecksum string) error { slog.Info("downloading update", "url", downloadURL) @@ -40,18 +44,18 @@ func selfUpdate(downloadURL, expectedChecksum string) error { } // Write to temp file next to current binary - currentExe, err := os.Executable() + currentExe, err := osExecutable() if err != nil { return fmt.Errorf("get executable path: %w", err) } tempPath := currentExe + ".update" - if err := os.WriteFile(tempPath, body, 0755); err != nil { + if err := osWriteFile(tempPath, body, 0755); err != nil { return fmt.Errorf("write temp: %w", err) } // Replace current binary - if err := os.Rename(tempPath, currentExe); err != nil { + if err := osRename(tempPath, currentExe); err != nil { _ = os.Remove(tempPath) return fmt.Errorf("rename: %w", err) } diff --git a/update_test.go b/update_test.go index b90b5b2..808dd68 100644 --- a/update_test.go +++ b/update_test.go @@ -5,6 +5,8 @@ import ( "fmt" "net/http" "net/http/httptest" + "os" + "strings" "testing" ) @@ -27,6 +29,116 @@ func TestSelfUpdateChecksumMismatch(t *testing.T) { } } +func TestSelfUpdate404(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + err := selfUpdate(srv.URL+"/missing", "") + 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.NewServer(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() + + err := selfUpdate(srv.URL, "") + // 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") + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("binary data")) + })) + defer srv.Close() + + err := selfUpdate(srv.URL, "") + 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 TestSelfUpdateWriteFileError(t *testing.T) { + origExe := osExecutable + origWrite := osWriteFile + defer func() { + osExecutable = origExe + osWriteFile = origWrite + }() + osExecutable = func() (string, error) { return "/tmp/test-agent", nil } + osWriteFile = func(name string, data []byte, perm os.FileMode) error { + return fmt.Errorf("disk full") + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("binary data")) + })) + defer srv.Close() + + err := selfUpdate(srv.URL, "") + if err == nil { + t.Error("expected write file error") + } + if !strings.Contains(err.Error(), "write temp") { + t.Errorf("expected 'write temp' in error, got: %v", err) + } +} + +func TestSelfUpdateRenameError(t *testing.T) { + origExe := osExecutable + origWrite := osWriteFile + origRename := osRename + defer func() { + osExecutable = origExe + osWriteFile = origWrite + osRename = origRename + }() + osExecutable = func() (string, error) { return "/tmp/test-agent", nil } + osWriteFile = func(name string, data []byte, perm os.FileMode) error { return nil } + osRename = func(oldpath, newpath string) error { + return fmt.Errorf("permission denied") + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("binary data")) + })) + defer srv.Close() + + err := selfUpdate(srv.URL, "") + 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)) diff --git a/websocket.go b/websocket.go index 22fa285..1b8969a 100644 --- a/websocket.go +++ b/websocket.go @@ -27,6 +27,12 @@ type WSConn struct { mu sync.Mutex // serializes writes } +var randRead = rand.Read +var netDial = net.Dial +var tlsDial = func(network, addr string) (net.Conn, error) { + return tls.Dial(network, addr, &tls.Config{MinVersion: tls.VersionTLS12}) +} + // WSDial connects to a WebSocket endpoint and performs the HTTP upgrade handshake. func WSDial(rawURL string) (*WSConn, error) { u, err := url.Parse(rawURL) @@ -46,9 +52,9 @@ func WSDial(rawURL string) (*WSConn, error) { var conn net.Conn if useTLS { - conn, err = tls.Dial("tcp", host, &tls.Config{MinVersion: tls.VersionTLS12}) + conn, err = tlsDial("tcp", host) } else { - conn, err = net.Dial("tcp", host) + conn, err = netDial("tcp", host) } if err != nil { return nil, fmt.Errorf("dial %s: %w", host, err) @@ -56,7 +62,7 @@ func WSDial(rawURL string) (*WSConn, error) { // Generate random key for Sec-WebSocket-Key keyBytes := make([]byte, 16) - if _, err := rand.Read(keyBytes); err != nil { + if _, err := randRead(keyBytes); err != nil { _ = conn.Close() return nil, fmt.Errorf("generate key: %w", err) } diff --git a/websocket_test.go b/websocket_test.go index 6825afd..a955975 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -3,11 +3,15 @@ package main import ( "bytes" "encoding/binary" + "fmt" + "io" + "net" + "net/http" + "strings" "testing" ) func TestWriteFrameMasked(t *testing.T) { - // Verify that writeFrame produces a valid masked client frame var buf bytes.Buffer ws := &WSConn{conn: &nopCloser{readWriter: &buf}} @@ -18,21 +22,17 @@ func TestWriteFrameMasked(t *testing.T) { frame := buf.Bytes() - // First byte: FIN + opcode - if frame[0] != 0x81 { // 0x80 (FIN) | 0x01 (text) + if frame[0] != 0x81 { t.Errorf("first byte: got %#x, want 0x81", frame[0]) } - // Second byte: MASK + length - if frame[1] != 0x85 { // 0x80 (mask) | 5 (length) + if frame[1] != 0x85 { t.Errorf("second byte: got %#x, want 0x85", frame[1]) } - // Mask key is bytes 2-5 maskKey := frame[2:6] maskedPayload := frame[6:] - // Unmask and verify for i := range maskedPayload { maskedPayload[i] ^= maskKey[i%4] } @@ -41,12 +41,11 @@ func TestWriteFrameMasked(t *testing.T) { } } -func TestWriteFrameExtendedLength(t *testing.T) { - // Test 16-bit extended length (126-65535 bytes) +func TestWriteFrame16BitLength(t *testing.T) { var buf bytes.Buffer ws := &WSConn{conn: &nopCloser{readWriter: &buf}} - payload := make([]byte, 300) // > 125, uses 2-byte extended + payload := make([]byte, 300) if err := ws.writeFrame(opBinary, payload); err != nil { t.Fatal(err) } @@ -61,8 +60,26 @@ func TestWriteFrameExtendedLength(t *testing.T) { } } +func TestWriteFrame64BitLength(t *testing.T) { + var buf bytes.Buffer + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + + payload := make([]byte, 70000) // > 65535, uses 8-byte extended + if err := ws.writeFrame(opBinary, payload); err != nil { + t.Fatal(err) + } + + frame := buf.Bytes() + if frame[1]&0x7F != 127 { + t.Errorf("expected 127 length marker, got %d", frame[1]&0x7F) + } + extLen := binary.BigEndian.Uint64(frame[2:10]) + if extLen != 70000 { + t.Errorf("extended length: got %d, want 70000", extLen) + } +} + func TestReadFrame(t *testing.T) { - // Build an unmasked server frame var buf bytes.Buffer payload := []byte("world") buf.WriteByte(0x81) // FIN + text @@ -82,8 +99,105 @@ func TestReadFrame(t *testing.T) { } } +func TestReadFrame16BitLength(t *testing.T) { + var buf bytes.Buffer + payload := make([]byte, 300) + for i := range payload { + payload[i] = byte(i % 256) + } + buf.WriteByte(0x82) // FIN + binary + buf.WriteByte(126) // 16-bit extended length + var extLen [2]byte + binary.BigEndian.PutUint16(extLen[:], 300) + buf.Write(extLen[:]) + buf.Write(payload) + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + opcode, data, err := ws.readFrame() + if err != nil { + t.Fatal(err) + } + if opcode != opBinary { + t.Errorf("opcode: got %d, want %d", opcode, opBinary) + } + if len(data) != 300 { + t.Errorf("data length: got %d, want 300", len(data)) + } +} + +func TestReadFrame64BitLength(t *testing.T) { + var buf bytes.Buffer + payload := make([]byte, 70000) + buf.WriteByte(0x82) // FIN + binary + buf.WriteByte(127) // 64-bit extended length + var extLen [8]byte + binary.BigEndian.PutUint64(extLen[:], 70000) + buf.Write(extLen[:]) + buf.Write(payload) + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + opcode, data, err := ws.readFrame() + if err != nil { + t.Fatal(err) + } + if opcode != opBinary { + t.Errorf("opcode: got %d, want %d", opcode, opBinary) + } + if len(data) != 70000 { + t.Errorf("data length: got %d, want 70000", len(data)) + } +} + +func TestReadFrameMasked(t *testing.T) { + var buf bytes.Buffer + payload := []byte("test") + maskKey := [4]byte{0x12, 0x34, 0x56, 0x78} + masked := make([]byte, len(payload)) + for i := range payload { + masked[i] = payload[i] ^ maskKey[i%4] + } + + buf.WriteByte(0x81) // FIN + text + buf.WriteByte(0x80 | byte(len(payload))) // masked + length + buf.Write(maskKey[:]) + buf.Write(masked) + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + opcode, data, err := ws.readFrame() + if err != nil { + t.Fatal(err) + } + if opcode != opText { + t.Errorf("opcode: got %d, want %d", opcode, opText) + } + if string(data) != "test" { + t.Errorf("data: got %q, want %q", data, "test") + } +} + +func TestReadMessageClose(t *testing.T) { + var buf bytes.Buffer + // Close frame + buf.WriteByte(0x80 | byte(opClose)) + buf.WriteByte(0) // no payload + + rw := &captureWriter{Reader: &buf} + ws := &WSConn{conn: &nopCloser{readWriter: rw}} + + _, opcode, err := ws.ReadMessage() + if err != io.EOF { + t.Errorf("expected io.EOF, got %v", err) + } + if opcode != opClose { + t.Errorf("opcode: got %d, want %d", opcode, opClose) + } + // Verify close reply was sent + if len(rw.written) == 0 { + t.Error("expected close reply frame to be written") + } +} + func TestReadFramePingPong(t *testing.T) { - // Server sends a ping, ReadMessage should auto-respond with pong and continue var buf bytes.Buffer // Ping frame @@ -107,12 +221,387 @@ func TestReadFramePingPong(t *testing.T) { t.Errorf("got %q, want %q", data, "data") } - // Verify pong was written if len(rw.written) == 0 { t.Error("expected pong frame to be written") } } +func TestReadFrame16BitLengthError(t *testing.T) { + var buf bytes.Buffer + buf.WriteByte(0x81) // FIN + text + buf.WriteByte(126) // 16-bit length indicator, but no length bytes follow + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + _, _, err := ws.readFrame() + if err == nil { + t.Error("expected error for truncated 16-bit length") + } +} + +func TestReadFrame64BitLengthError(t *testing.T) { + var buf bytes.Buffer + buf.WriteByte(0x81) // FIN + text + buf.WriteByte(127) // 64-bit length indicator, but no length bytes follow + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + _, _, err := ws.readFrame() + if err == nil { + t.Error("expected error for truncated 64-bit length") + } +} + +func TestReadFrameMaskKeyError(t *testing.T) { + var buf bytes.Buffer + buf.WriteByte(0x81) // FIN + text + buf.WriteByte(0x80 | 0x01) // masked + length 1, but no mask key or payload + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + _, _, err := ws.readFrame() + if err == nil { + t.Error("expected error for missing mask key") + } +} + +func TestReadFramePayloadError(t *testing.T) { + var buf bytes.Buffer + buf.WriteByte(0x81) // FIN + text + buf.WriteByte(5) // length 5, but no payload + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + _, _, err := ws.readFrame() + if err == nil { + t.Error("expected error for truncated payload") + } +} + +func TestWriteFrameHeaderError(t *testing.T) { + ws := &WSConn{conn: &failOnWriteBuffer{Buffer: bytes.NewBuffer(nil)}} + err := ws.writeFrame(opText, []byte("hello")) + if err == nil { + t.Error("expected write error") + } +} + +func TestWriteFramePayloadError(t *testing.T) { + // Writer that succeeds for header but fails for payload + ws := &WSConn{conn: &failOnNthWrite{failAfter: 1}} + err := ws.writeFrame(opText, []byte("hello")) + if err == nil { + t.Error("expected write error on payload") + } +} + +// failOnNthWrite fails after N successful writes. +type failOnNthWrite struct { + failAfter int + count int +} + +func (f *failOnNthWrite) Read(p []byte) (int, error) { return 0, io.EOF } +func (f *failOnNthWrite) Write(p []byte) (int, error) { + f.count++ + if f.count > f.failAfter { + return 0, fmt.Errorf("write failed") + } + return len(p), nil +} +func (f *failOnNthWrite) Close() error { return nil } + +func TestReadMessageError(t *testing.T) { + // Empty buffer causes immediate EOF on readFrame + buf := bytes.NewBuffer(nil) + ws := &WSConn{conn: &nopCloser{readWriter: buf}} + _, _, err := ws.ReadMessage() + if err == nil { + t.Error("expected error for empty buffer") + } +} + +func TestReadMessagePongError(t *testing.T) { + // Ping frame followed by write failure for pong + var buf bytes.Buffer + buf.WriteByte(0x80 | byte(opPing)) + buf.WriteByte(0) + + ws := &WSConn{conn: &nopCloser{readWriter: &failOnWriteBuffer{Buffer: &buf}}} + _, _, err := ws.ReadMessage() + if err == nil { + t.Error("expected pong write error") + } +} + +func TestReadFrameEmptyPayload(t *testing.T) { + var buf bytes.Buffer + buf.WriteByte(0x82) // FIN + binary + buf.WriteByte(0) // zero length + + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + opcode, data, err := ws.readFrame() + if err != nil { + t.Fatal(err) + } + if opcode != opBinary { + t.Errorf("opcode: got %d, want %d", opcode, opBinary) + } + if len(data) != 0 { + t.Errorf("expected empty payload, got %d bytes", len(data)) + } +} + +func TestWriteFrameEmptyPayload(t *testing.T) { + var buf bytes.Buffer + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + + if err := ws.writeFrame(opClose, nil); err != nil { + t.Fatal(err) + } + + frame := buf.Bytes() + if frame[0]&0x0F != opClose { + t.Errorf("expected close opcode, got %d", frame[0]&0x0F) + } + // Length should be 0 (masked) + if frame[1]&0x7F != 0 { + t.Errorf("expected 0 length, got %d", frame[1]&0x7F) + } +} + +// failOnWriteBuffer reads from Buffer but fails on Write. +type failOnWriteBuffer struct { + *bytes.Buffer +} + +func (f *failOnWriteBuffer) Write(p []byte) (int, error) { + return 0, fmt.Errorf("write failed") +} +func (f *failOnWriteBuffer) Close() error { return nil } + +func TestWriteText(t *testing.T) { + var buf bytes.Buffer + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + + if err := ws.WriteText([]byte("hello")); err != nil { + t.Fatal(err) + } + + // Verify it wrote a text frame (opcode 1) + frame := buf.Bytes() + if frame[0]&0x0F != opText { + t.Errorf("expected text opcode, got %d", frame[0]&0x0F) + } +} + +func TestClose(t *testing.T) { + var buf bytes.Buffer + ws := &WSConn{conn: &nopCloser{readWriter: &buf}} + + if err := ws.Close(); err != nil { + t.Fatal(err) + } + + // Verify it wrote a close frame (opcode 8) + frame := buf.Bytes() + if len(frame) == 0 { + t.Fatal("expected close frame to be written") + } + if frame[0]&0x0F != opClose { + t.Errorf("expected close opcode, got %d", frame[0]&0x0F) + } +} + +func TestWSDial(t *testing.T) { + // Start a test TCP server that responds with a valid WebSocket upgrade + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + + buf := make([]byte, 4096) + n, _ := conn.Read(buf) + _ = string(buf[:n]) // Read the request + + resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n" + conn.Write([]byte(resp)) + + // Keep connection open briefly for the test + buf2 := make([]byte, 1) + conn.Read(buf2) + }() + + addr := ln.Addr().String() + ws, err := WSDial("ws://" + addr + "/socket") + if err != nil { + t.Fatal(err) + } + _ = ws.Close() +} + +func TestWSDialHandshakeFailed(t *testing.T) { + // Start a test HTTP server that returns 403 + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, err := ln.Accept() + if err != nil { + return + } + defer conn.Close() + + buf := make([]byte, 4096) + conn.Read(buf) + + resp := "HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\n\r\n" + conn.Write([]byte(resp)) + }() + + addr := ln.Addr().String() + _, err = WSDial("ws://" + addr + "/socket") + if err == nil { + t.Error("expected handshake error") + } + if !strings.Contains(err.Error(), "handshake failed") { + t.Errorf("expected 'handshake failed' in error, got: %v", err) + } +} + +func TestWSDialReadHandshakeError(t *testing.T) { + // Server accepts connection but immediately closes without sending response + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, _ := ln.Accept() + if conn != nil { + // Read the request then close immediately + buf := make([]byte, 4096) + conn.Read(buf) + conn.Close() + } + }() + + addr := ln.Addr().String() + _, err = WSDial("ws://" + addr + "/socket") + if err == nil { + t.Error("expected read handshake error") + } +} + +func TestWSDialGenerateKeyError(t *testing.T) { + origRandRead := randRead + defer func() { randRead = origRandRead }() + randRead = func(b []byte) (int, error) { + return 0, fmt.Errorf("entropy exhausted") + } + + // Start a TCP server that accepts connections + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + go func() { + conn, _ := ln.Accept() + if conn != nil { + conn.Close() + } + }() + + addr := ln.Addr().String() + _, err = WSDial("ws://" + addr + "/socket") + if err == nil { + t.Error("expected generate key error") + } + if !strings.Contains(err.Error(), "generate key") { + t.Errorf("expected 'generate key' in error, got: %v", err) + } +} + +func TestWSDialWriteHandshakeError(t *testing.T) { + origDial := netDial + defer func() { netDial = origDial }() + + netDial = func(network, addr string) (net.Conn, error) { + // Return a connection whose Write always fails + client, _ := net.Pipe() + client.Close() // Close immediately so Write fails + return client, nil + } + + _, err := WSDial("ws://127.0.0.1:9999/socket") + if err == nil { + t.Error("expected write handshake error") + } + if !strings.Contains(err.Error(), "write handshake") { + t.Errorf("expected 'write handshake' in error, got: %v", err) + } +} + +func TestWSDialBadURL(t *testing.T) { + _, err := WSDial("://bad url") + if err == nil { + t.Error("expected parse error for bad URL") + } +} + +func TestWSDialDefaultPorts(t *testing.T) { + // Test that ws:// defaults to port 80 — will fail to connect but verifies URL parsing + _, err := WSDial("ws://127.0.0.1/path") + if err == nil { + t.Error("expected connection error (nothing on port 80)") + } + + // Test that wss:// defaults to port 443 + _, err = WSDial("wss://127.0.0.1/path") + if err == nil { + t.Error("expected connection error (nothing on port 443)") + } +} + +func TestWSDialRealHTTPServer(t *testing.T) { + // Test with a real HTTP server that sends proper 101 upgrade + mux := http.NewServeMux() + mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { + hj, ok := w.(http.Hijacker) + if !ok { + http.Error(w, "hijack not supported", 500) + return + } + conn, brw, _ := hj.Hijack() + defer conn.Close() + brw.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n") + brw.Flush() + // Keep alive briefly + buf := make([]byte, 1) + conn.Read(buf) + }) + ln, _ := net.Listen("tcp", "127.0.0.1:0") + defer ln.Close() + go http.Serve(ln, mux) + + addr := ln.Addr().String() + ws, err := WSDial("ws://" + addr + "/ws") + if err != nil { + t.Fatal(err) + } + _ = ws.Close() +} + // nopCloser wraps a ReadWriter with a no-op Close. type nopCloser struct { readWriter interface {