graham/towerops-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

increase test coverage and improve logger output

Browse source
This commit is contained in:
Graham McIntire 2026-02-11 10:42:19 -06:00
parent a757bd7615
commit a36a2b7fd7
No known key found for this signature in database
17 changed files with 2910 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -2,6 +2,7 @@
towerops-agent
*.test
cover.out
cover.html
# Database files
*.db

7
agent.go
View file

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

349
agent_test.go
View file

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

118
log.go Normal file
View file

@ -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,
}
}

159
log_test.go Normal file
View file

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

2
main.go
View file

@ -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)")

22
main_test.go
View file

@ -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 {

6
mikrotik.go
View file

@ -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,

548
mikrotik_test.go
View file

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

43
ping_test.go
View file

@ -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) {

501
snmp_test.go
View file

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

5
ssh.go
View file

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

528
ssh_test.go Normal file
View file

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

10
update.go
View file

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

112
update_test.go
View file

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

12
websocket.go
View file

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

515
websocket_test.go
View file

@ -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 {