diff --git a/cmd/jobsearch/main.go b/cmd/jobsearch/main.go index 59619eb..adc5949 100644 --- a/cmd/jobsearch/main.go +++ b/cmd/jobsearch/main.go @@ -28,6 +28,7 @@ import ( "git.avsdev.uk/Go/job-search/internal/cli/fetch" "git.avsdev.uk/Go/job-search/internal/cli/migrate" "git.avsdev.uk/Go/job-search/internal/cli/version" + "git.avsdev.uk/Go/job-search/internal/cli/wsserve" "github.com/spf13/cobra" ) @@ -75,6 +76,7 @@ opportunities. config.NewCommand(opts), fetch.NewCommand(opts), migrate.NewCommand(opts), + wsserve.NewCommand(opts), ) return root diff --git a/cmd/jobsearch/main_test.go b/cmd/jobsearch/main_test.go index f066e0d..bde0dce 100644 --- a/cmd/jobsearch/main_test.go +++ b/cmd/jobsearch/main_test.go @@ -16,7 +16,7 @@ func TestCommandTree(t *testing.T) { for _, c := range root.Commands() { names = append(names, c.Name()) } - for _, want := range []string{"config", "db", "fetch", "migrate", "version"} { + for _, want := range []string{"config", "fetch", "migrate", "version", "wsserve"} { if !slices.Contains(names, want) { t.Errorf("job-search has no %q command; has %v", want, names) } diff --git a/go.mod b/go.mod index 4b0a20c..f795ee3 100644 --- a/go.mod +++ b/go.mod @@ -3,14 +3,20 @@ module git.avsdev.uk/Go/job-search go 1.26.8 require ( + github.com/coder/websocket v1.8.15 + github.com/dustin/go-humanize v1.0.1 + github.com/go-chi/chi/v5 v5.4.0 github.com/jmoiron/sqlx v1.4.0 github.com/joho/godotenv v1.5.1 + github.com/klauspost/compress v1.20.3 github.com/lmittmann/tint v1.2.0 github.com/mattn/go-sqlite3 v1.14.52 + github.com/pires/go-proxyproto v0.15.0 github.com/pressly/goose/v3 v3.28.0 github.com/spf13/cobra v1.10.2 github.com/spf13/viper v1.21.0 go.yaml.in/yaml/v3 v3.0.5 + golang.org/x/time v0.16.0 resty.dev/v3 v3.0.0-rc.4 ) diff --git a/go.sum b/go.sum index cf7653a..cbc68f8 100644 --- a/go.sum +++ b/go.sum @@ -1,6 +1,8 @@ filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo= filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc= +github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA= +github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= @@ -8,6 +10,9 @@ github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHk github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.10.1 h1:b0/UzAf9yR5rhf3RPm9gf3ehBPpf0oZKIjtpKrx59Ho= github.com/fsnotify/fsnotify v1.10.1/go.mod h1:TLheqan6HD6GBK6PrDWyDPBaEV8LspOxvPSjC+bVfgo= +github.com/go-chi/chi/v5 v5.3.2/go.mod h1:R+tYY2hNuVUUjxoPtqUdgBqevM9s9njzkTLutVsOCto= +github.com/go-chi/chi/v5 v5.4.0 h1:+ufzvFTGKyin26Or5DLrlIDq4ygIDBVcj7oVNeeaYFE= +github.com/go-chi/chi/v5 v5.4.0/go.mod h1:d1IG7GtKPSjPyBP365PH8sAGVHk3qfB/IhjFNRH1wWM= github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg= github.com/go-sql-driver/mysql v1.10.0 h1:Q+1LV8DkHJvSYAdR83XzuhDaTykuDx0l6fkXxoWCWfw= github.com/go-sql-driver/mysql v1.10.0/go.mod h1:M+cqaI7+xxXGG9swrdeUIoPG3Y3KCkF0pZej+SK+nWk= @@ -23,6 +28,9 @@ github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o= github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY= github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0= github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= +github.com/klauspost/compress v1.20.1/go.mod h1:LUdAzn7YLVvxLpc7y3V1m40wESHTgc1422pwwBSKYuI= +github.com/klauspost/compress v1.20.3 h1:zWNoFqq/lSuYQRim3dUXlXeA0IqwuZGicdG5ZTfY5j4= +github.com/klauspost/compress v1.20.3/go.mod h1:LUdAzn7YLVvxLpc7y3V1m40wESHTgc1422pwwBSKYuI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -42,6 +50,8 @@ github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOF github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pires/go-proxyproto v0.15.0 h1:dTshmNbFm/D+0+sbrxUuddPOZ5Y0B7c5NhtsBkm6LqI= +github.com/pires/go-proxyproto v0.15.0/go.mod h1:OXsCrKwrK2tXS9YrI5tkHx5xaQlO8FH3lFW76orFh24= github.com/pressly/goose/v3 v3.28.0 h1:D2M+iL31GmpZxSHOhX8mqyqAT3CXnokUmm0eKoSP+Vc= github.com/pressly/goose/v3 v3.28.0/go.mod h1:v26MOuB8bL3kzzrt3Vqhb3R0PRVsl8hFQKdrht/L6Rk= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= @@ -83,6 +93,9 @@ golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +golang.org/x/time v0.16.0 h1:vMb6ptszcQMkcwiRTAuNNU50gom6++Q/6gY2hDM6VDE= +golang.org/x/time v0.16.0/go.mod h1:rVKOqvZeKvrDKTQiAHJ7wmwP0RzleSphoEA9RcdLA0s= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= modernc.org/libc v1.75.6 h1:yKk8qo+Di4gkmvRboK8ocCqH22FiUCR6jRy2OwtCRus= modernc.org/libc v1.75.6/go.mod h1:bO5o2ztHxBb2rjz0PgdHN0sSMw57CgxGFLZ3Qd/QpVQ= diff --git a/internal/cli/wsserve/command.go b/internal/cli/wsserve/command.go new file mode 100644 index 0000000..f126815 --- /dev/null +++ b/internal/cli/wsserve/command.go @@ -0,0 +1,89 @@ +package wsserve + +import ( + "context" + "fmt" + "log/slog" + "os" + "os/signal" + "strconv" + "strings" + "syscall" + + "git.avsdev.uk/Go/job-search/internal/cli" + "git.avsdev.uk/Go/job-search/internal/websocket" + "git.avsdev.uk/Go/job-search/internal/websocket/tcpserver" + "github.com/spf13/cobra" +) + +func NewCommand(opts *cli.Options) *cobra.Command { + cmd := &cobra.Command{ + Use: "wsserve", + Short: "Run JobSearch TCP Websocket server", + PersistentPreRun: func(cmd *cobra.Command, args []string) { + slog.Debug("wsserve.PersistentPreRun") + }, + PersistentPostRun: func(cmd *cobra.Command, args []string) { + slog.Debug("wsserve.PersistentPostRun") + }, + RunE: func(cmd *cobra.Command, args []string) error { + slog.Debug("wsserve.Run") + + addrOverride, _ := cmd.Flags().GetString("listen") + hostOverride, _ := cmd.Flags().GetString("host") + portOverride, _ := cmd.Flags().GetInt("port") + + if addrOverride != "" { + listenParts := strings.Split(addrOverride, ":") + if len(listenParts) != 2 { + return fmt.Errorf( + "--listen should be in the format of 127.0.0.1:9393", + ) + } + i, err := strconv.Atoi(listenParts[1]) + if err != nil { + return fmt.Errorf("Failed to parse --listen port: %w", err) + } + hostOverride = listenParts[0] + portOverride = i + } + + conf, err := opts.RunSetup(cmd, true, true) + if err != nil { + return fmt.Errorf("no configuration loaded (%w)", err) + } + + if portOverride != 0 { + conf.Server.Port = portOverride + } + if hostOverride != "" { + conf.Server.Host = hostOverride + } + + // ctx is cancelled on the first SIGINT (Ctrl+C) or SIGTERM. + ctx, stop := signal.NotifyContext( + context.Background(), os.Interrupt, syscall.SIGTERM, + ) + defer stop() + + // After the first signal, restore default signal handling so a second + // Ctrl+C force-kills the process if graceful shutdown hangs. + go func() { + <-ctx.Done() + stop() + }() + + hub := websocket.NewHub(conf.Server.WebSocket, slog.Default()) + svr := tcpserver.New(&conf.Server, hub, slog.Default()) + + return svr.Run(ctx) + }, + } + cmd.Flags().String("listen", "", "listen address (default from config, 127.0.0.1:9393)") + cmd.Flags().String("host", "", "listen address (default from config, 127.0.0.1)") + cmd.Flags().Int("port", 0, "listen address (default from config, 9393)") + cmd.Flags().Bool("foreground", false, "run in the foreground") + cmd.MarkFlagsMutuallyExclusive("listen", "host") + cmd.MarkFlagsMutuallyExclusive("listen", "port") + return cmd +} diff --git a/internal/config/README.md b/internal/config/README.md index 58b4bf0..41726b5 100644 --- a/internal/config/README.md +++ b/internal/config/README.md @@ -40,11 +40,26 @@ is kept. | `sources.find_a_tender.api_root` | string | `https://www.find-tender.service.gov.uk/api/1.0` | URL of the Find a Tender API | | `sources.fts_mock.api_root` | string | | URL of the mock FTS server (see `internal/mock`) | | `sources.sleeper.block_shutdown` | bool | `false` | Test source: if true, its sleep ignores shutdown | +| `server.host` | string | `127.0.0.1` | Address to listen on | +| `server.port` | int | `9393` | Port to listen on | +| `server.timeouts.shutdown` | duration | `15s` | How long shutdown waits for clients to close | +| `server.rate_limits.idle_ttl` | duration | `10m` | How long a client's rate limiter is kept after its last request | +| `server.websocket.enable_compression` | bool | `false` | permessage-deflate | +| `server.websocket.max_message_size` | byte size | `32KiB` | Largest message read from a client; `0B` for 32KiB | +| `server.websocket.connection_limit.*` | | | A rate limit rule for new connections from each client IP, WebSocket and plain TCP alike; limiters are kept for `server.rate_limits.idle_ttl` | +| `server.websocket.message_rate_limit.*` | | | A rate limit rule for messages from each connection | +| `server.websocket.write_timeout` | duration | `10s` | `0s` for none | +| `server.websocket.ping_interval` | duration | `30s` | | +| `server.websocket.pong_timeout` | duration | `10s` | `0s` for none | The current sources are `find_a_tender`, `fts_mock` and `sleeper`. See "Addition of a new opportunities source" in the top level README for how to add a new source. +A rate limit rule (`JSCRateLimitRule`) has `enabled`, `per_second` and `burst`. +It is used under `server.websocket.connection_limit` and +`server.websocket.message_rate_limit`. + Example `config.yaml`: ```yaml logging: @@ -79,6 +94,9 @@ per key: `time.Duration` or a duration string in `SetValue`, and are given out as strings (`"10m0s"`) by `Map`, `Flatten` and `ValueOf`, and so by `JSON`, `YAML` and `String` too. +- `util.ByteSize` fields work the same way, with size strings such as `"32KiB"` + or `"1MB"` (KB, MB, ... are powers of 1000; KiB, MiB, ... powers of 1024), + and are given out exactly in the largest unit that divides them (`"1MB"`). - A section can reject values by implementing `validate(key, value) error` on its pointer type. `JobSearchConfigLogging` does this for `level`. @@ -138,6 +156,9 @@ Nothing in `section.go` needs changing. | `conf_logging.go` | `logging` section and log level helpers | | `conf_sources.go` | `sources` section and the shared source fields | | `conf_source_*.go` | One file per source section | +| `conf_server.go` | `server` section | +| `conf_server_*.go` | One file per `server` sub-section (`rate_limits`, `timeouts`, `websocket`) | +| `conf_ratelimitrule.go` | `JSCRateLimitRule`, the rate limit rule shared by several keys | | `print.go`, `debug.go` | Printing and debug logging of the configuration | diff --git a/internal/config/conf_ratelimitrule.go b/internal/config/conf_ratelimitrule.go new file mode 100644 index 0000000..b33f8b6 --- /dev/null +++ b/internal/config/conf_ratelimitrule.go @@ -0,0 +1,29 @@ +package config + +import "reflect" + +// JSCRateLimitRule configures one rate limit. The type is used under several +// keys so it has no JSCK* constant and its Flatten keys are not prefixed. +type JSCRateLimitRule struct { + Enabled bool `cfg:"enabled" default:"false"` + PerSecond int `cfg:"per_second" default:"50"` + Burst int `cfg:"burst" default:"100"` +} + +// ---------- Section methods (see Section) ---------- + +func (c JSCRateLimitRule) Map() map[string]any { return mapOf(c) } + +func (c JSCRateLimitRule) Flatten() map[string]any { return flattenOf("", c) } + +func (c JSCRateLimitRule) TypeOf(key string) (reflect.Kind, error) { + return typeOf(c, key) +} + +func (c JSCRateLimitRule) ValueOf(key string) (any, error) { + return valueOf(c, key) +} + +func (c *JSCRateLimitRule) SetValue(key string, value any) error { + return setValue(c, key, value) +} diff --git a/internal/config/conf_sections_test.go b/internal/config/conf_sections_test.go index 601deb4..c9593ad 100644 --- a/internal/config/conf_sections_test.go +++ b/internal/config/conf_sections_test.go @@ -50,10 +50,10 @@ func TestDefaults(t *testing.T) { if conf.Sources.FTSMock.APIRoot != "" { t.Errorf("FTSMock.APIRoot = %q; want empty", conf.Sources.FTSMock.APIRoot) } - for name, src := range map[string]JobSearchConfigSource{ - "find_a_tender": conf.Sources.FindATender.JobSearchConfigSource, - "fts_mock": conf.Sources.FTSMock.JobSearchConfigSource, - "sleeper": conf.Sources.Sleeper.JobSearchConfigSource, + for name, src := range map[string]JSCSourcesSource{ + "find_a_tender": conf.Sources.FindATender.JSCSourcesSource, + "fts_mock": conf.Sources.FTSMock.JSCSourcesSource, + "sleeper": conf.Sources.Sleeper.JSCSourcesSource, } { if src.Enabled || src.Interval != 10*time.Minute { t.Errorf("%s source = %+v; want disabled, 10m", name, src) @@ -150,7 +150,7 @@ func TestConfigLevelToLogLevel(t *testing.T) { } func TestSourceSection(t *testing.T) { - c := JobSearchConfigSource{Enabled: true, Interval: time.Minute} + c := JSCSourcesSource{Enabled: true, Interval: time.Minute} if got, want := c.Keys(), []string{"enabled", "interval"}; !reflect.DeepEqual(got, want) { t.Errorf("Keys = %v; want %v", got, want) @@ -352,6 +352,7 @@ func TestSerialisedRoundTrip(t *testing.T) { orig.Logging.Level = "info" orig.Sources.FTSMock.Enabled = true orig.Sources.FTSMock.APIRoot = "http://mock" + orig.Server.WebSocket.MaxMessageSize = 64 << 10 for _, format := range []string{"json", "yaml"} { t.Run(format, func(t *testing.T) { resetViper(t) diff --git a/internal/config/conf_server.go b/internal/config/conf_server.go new file mode 100644 index 0000000..9f54df0 --- /dev/null +++ b/internal/config/conf_server.go @@ -0,0 +1,37 @@ +package config + +import ( + "reflect" +) + +// JSCKServer is the server section's configuration key. +const JSCKServer string = "server" + +// JSCServer configures the outward server. +type JSCServer struct { + Host string `cfg:"host" default:"127.0.0.1"` + Port int `cfg:"port" default:"9393"` + Timeouts JSCServerTimeouts `cfg:"timeouts"` + RateLimits JSCServerRateLimits `cfg:"rate_limits"` + WebSocket JSCServerWebSocket `cfg:"websocket"` +} + +// ---------- Section methods (see Section) ---------- + +func (c JSCServer) Map() map[string]any { return mapOf(c) } + +func (c JSCServer) Flatten() map[string]any { + return flattenOf(JSCKServer, c) +} + +func (c JSCServer) TypeOf(key string) (reflect.Kind, error) { + return typeOf(c, key) +} + +func (c JSCServer) ValueOf(key string) (any, error) { + return valueOf(c, key) +} + +func (c *JSCServer) SetValue(key string, value any) error { + return setValue(c, key, value) +} diff --git a/internal/config/conf_server_ratelimits.go b/internal/config/conf_server_ratelimits.go new file mode 100644 index 0000000..58aa87e --- /dev/null +++ b/internal/config/conf_server_ratelimits.go @@ -0,0 +1,35 @@ +package config + +import ( + "reflect" + "time" +) + +// JSCKServerRateLimits is the server rate limits section's configuration key. +const JSCKServerRateLimits string = "rate_limits" + +// JSCServerRateLimits configures the server's rate limiters. IdleTTL is how +// long a client's limiter is kept after its last request. +type JSCServerRateLimits struct { + IdleTTL time.Duration `cfg:"idle_ttl" default:"10m"` +} + +// ---------- Section methods (see Section) ---------- + +func (c JSCServerRateLimits) Map() map[string]any { return mapOf(c) } + +func (c JSCServerRateLimits) Flatten() map[string]any { + return flattenOf(JSCKServerRateLimits, c) +} + +func (c JSCServerRateLimits) TypeOf(key string) (reflect.Kind, error) { + return typeOf(c, key) +} + +func (c JSCServerRateLimits) ValueOf(key string) (any, error) { + return valueOf(c, key) +} + +func (c *JSCServerRateLimits) SetValue(key string, value any) error { + return setValue(c, key, value) +} diff --git a/internal/config/conf_server_timeouts.go b/internal/config/conf_server_timeouts.go new file mode 100644 index 0000000..a9be776 --- /dev/null +++ b/internal/config/conf_server_timeouts.go @@ -0,0 +1,35 @@ +package config + +import ( + "reflect" + "time" +) + +// JSCKServerTimeouts is the server timeouts section's configuration key. +const JSCKServerTimeouts string = "timeouts" + +// JSCServerTimeouts configures the server's timeouts. ShutdownTimeout is how +// long shutdown waits for clients to close. +type JSCServerTimeouts struct { + ShutdownTimeout time.Duration `cfg:"shutdown" default:"15s"` +} + +// ---------- Section methods (see Section) ---------- + +func (c JSCServerTimeouts) Map() map[string]any { return mapOf(c) } + +func (c JSCServerTimeouts) Flatten() map[string]any { + return flattenOf(JSCKServerTimeouts, c) +} + +func (c JSCServerTimeouts) TypeOf(key string) (reflect.Kind, error) { + return typeOf(c, key) +} + +func (c JSCServerTimeouts) ValueOf(key string) (any, error) { + return valueOf(c, key) +} + +func (c *JSCServerTimeouts) SetValue(key string, value any) error { + return setValue(c, key, value) +} diff --git a/internal/config/conf_server_websocket.go b/internal/config/conf_server_websocket.go new file mode 100644 index 0000000..ede8957 --- /dev/null +++ b/internal/config/conf_server_websocket.go @@ -0,0 +1,42 @@ +package config + +import ( + "reflect" + "time" + + "git.avsdev.uk/Go/job-search/internal/util" +) + +// JSCKServerWebSocket is the websocket section's configuration key. +const JSCKServerWebSocket string = "websocket" + +// JSCServerWebSocket configures the server's websocket connections. +type JSCServerWebSocket struct { + EnableCompression bool `cfg:"enable_compression" default:"false"` + MaxMessageSize util.ByteSize `cfg:"max_message_size" default:"32KiB"` + ConnectionLimit JSCRateLimitRule `cfg:"connection_limit"` + MessageRateLimit JSCRateLimitRule `cfg:"message_rate_limit"` + WriteTimeout time.Duration `cfg:"write_timeout" default:"10s"` + PingInterval time.Duration `cfg:"ping_interval" default:"30s"` + PongTimeout time.Duration `cfg:"pong_timeout" default:"10s"` +} + +// ---------- Section methods (see Section) ---------- + +func (c JSCServerWebSocket) Map() map[string]any { return mapOf(c) } + +func (c JSCServerWebSocket) Flatten() map[string]any { + return flattenOf(JSCKServerWebSocket, c) +} + +func (c JSCServerWebSocket) TypeOf(key string) (reflect.Kind, error) { + return typeOf(c, key) +} + +func (c JSCServerWebSocket) ValueOf(key string) (any, error) { + return valueOf(c, key) +} + +func (c *JSCServerWebSocket) SetValue(key string, value any) error { + return setValue(c, key, value) +} diff --git a/internal/config/conf_source_findatender.go b/internal/config/conf_source_findatender.go index 7cf038e..50bb9c5 100644 --- a/internal/config/conf_source_findatender.go +++ b/internal/config/conf_source_findatender.go @@ -8,7 +8,7 @@ const JSCKSourceFindATender string = "find_a_tender" // JSCSourceFindATender configures fetching from the Find a Tender // OCDS API at APIRoot. type JSCSourceFindATender struct { - JobSearchConfigSource + JSCSourcesSource APIRoot string `cfg:"api_root" default:"https://www.find-tender.service.gov.uk/api/1.0"` } diff --git a/internal/config/conf_source_ftsmock.go b/internal/config/conf_source_ftsmock.go index 93b146f..9759e14 100644 --- a/internal/config/conf_source_ftsmock.go +++ b/internal/config/conf_source_ftsmock.go @@ -8,7 +8,7 @@ const JSCKSourceFTSMock string = "fts_mock" // JSCSourceFTSMock configures fetching from a Find a Tender mock // server (cmd/mock-fts) at APIRoot. type JSCSourceFTSMock struct { - JobSearchConfigSource + JSCSourcesSource APIRoot string `cfg:"api_root"` } diff --git a/internal/config/conf_source_sleeper.go b/internal/config/conf_source_sleeper.go index 383d90e..2f9b308 100644 --- a/internal/config/conf_source_sleeper.go +++ b/internal/config/conf_source_sleeper.go @@ -9,7 +9,7 @@ const JSCKSourceSleeper string = "sleeper" // simulates work by sleeping. With BlockShutdown its work ignores // cancellation, to test shutdown waiting on a non-cooperative task. type JSCSourceSleeper struct { - JobSearchConfigSource + JSCSourcesSource BlockShutdown bool `cfg:"block_shutdown"` } diff --git a/internal/config/conf_sources.go b/internal/config/conf_sources.go index 1eebfee..0a82496 100644 --- a/internal/config/conf_sources.go +++ b/internal/config/conf_sources.go @@ -38,32 +38,32 @@ func (c *JSCSources) SetValue(key string, value any) error { // ========== Generic source ========== -// JobSearchConfigSource holds the settings every source has, and is embedded +// JSCSourcesSource holds the settings every source has, and is embedded // in each source's section. Interval is how often an enabled source fetches. -type JobSearchConfigSource struct { +type JSCSourcesSource struct { Enabled bool `cfg:"enabled"` Interval time.Duration `cfg:"interval" default:"10m"` } // ---------- Section methods (see Section) ---------- -func (c JobSearchConfigSource) Map() map[string]any { return mapOf(c) } +func (c JSCSourcesSource) Map() map[string]any { return mapOf(c) } -func (c JobSearchConfigSource) Flatten() map[string]any { +func (c JSCSourcesSource) Flatten() map[string]any { return flattenOf("", c) } -func (c JobSearchConfigSource) TypeOf(key string) (reflect.Kind, error) { +func (c JSCSourcesSource) TypeOf(key string) (reflect.Kind, error) { return typeOf(c, key) } -func (c JobSearchConfigSource) ValueOf(key string) (any, error) { +func (c JSCSourcesSource) ValueOf(key string) (any, error) { return valueOf(c, key) } -func (c *JobSearchConfigSource) SetValue(key string, value any) error { +func (c *JSCSourcesSource) SetValue(key string, value any) error { return setValue(c, key, value) } // Keys returns the keys directly in the section. -func (c *JobSearchConfigSource) Keys() []string { return keysOf(c) } +func (c *JSCSourcesSource) Keys() []string { return keysOf(c) } diff --git a/internal/config/config.go b/internal/config/config.go index 771e995..ed5ef56 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -24,9 +24,10 @@ const EtcConfigPath string = "/etc/job-search" // JobSearchConfig is the application's whole configuration, one field per // section. Each field's cfg tag must match the section's JSCK* constant. type JobSearchConfig struct { - Database JSCDatabase `cfg:"database"` - Logging JSCLogging `cfg:"logging"` - Sources JSCSources `cfg:"sources"` + Database JSCDatabase `cfg:"database"` + Logging JSCLogging `cfg:"logging"` + Server JSCServer `cfg:"server"` + Sources JSCSources `cfg:"sources"` } // LoadDefaultConfig sets Viper's defaults from the configuration's default @@ -88,6 +89,24 @@ func viperValue(key string, kind reflect.Kind) (value any, ok bool) { return viper.GetInt(key), true case reflect.Bool: return viper.GetBool(key), true + case reflect.Slice: + // A string, e.g. from an environment variable, is passed on for + // SetValue to split on commas; GetStringSlice would split it on + // spaces. A list of objects (a list of sections) is passed on as it + // is, as GetStringSlice would flatten it. + switch v := viper.Get(key).(type) { + case string: + return v, true + case []map[string]any: // as LoadToViper sets it + return v, true + case []any: + if len(v) > 0 { + if _, ok := v[0].(map[string]any); ok { + return v, true + } + } + } + return viper.GetStringSlice(key), true } return nil, false } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 40b99d3..13e8eb0 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1,13 +1,16 @@ package config import ( + "fmt" + "maps" "reflect" + "strings" "testing" ) // Kinds that LoadFromViper knows how to read. var viperKinds = map[reflect.Kind]bool{ - reflect.String: true, reflect.Int: true, reflect.Bool: true, + reflect.String: true, reflect.Int: true, reflect.Bool: true, reflect.Slice: true, } // Every key round-trips through TypeOf/ValueOf/SetValue and has a type that @@ -39,19 +42,40 @@ var sectionType = reflect.TypeOf((*Section)(nil)).Elem() // Every section, including embedded ones, implements Section, and its // Flatten keys (which use the section's JSCK* constant) line up with its cfg -// tag in the parent. Sections are found by walking JobSearchConfig, so a new -// section is checked without being listed anywhere. +// tag in the parent. A section type used under more than one key anywhere in +// the config, such as JSCRateLimitRule, has no constant, so its keys are +// unprefixed. Each section's Map, TypeOf, ValueOf and SetValue agree with the +// whole configuration's view of the same keys, and so do those of each item +// in a list of sections. Sections are found by walking JobSearchConfig, so a +// new section is checked without being listed anywhere. func TestSections(t *testing.T) { - var conf JobSearchConfig + conf := defaultConfig(t) if !reflect.PointerTo(reflect.TypeOf(conf)).Implements(sectionType) { t.Error("*JobSearchConfig does not implement Section") } all := conf.Flatten() - checkSections(t, all, "", reflect.ValueOf(&conf).Elem()) + uses := map[reflect.Type]int{} + countSections(uses, reflect.TypeOf(conf)) + checkSections(t, &conf, all, uses, "", reflect.ValueOf(&conf).Elem()) } -func checkSections(t *testing.T, all map[string]any, prefix string, section reflect.Value) { +// countSections counts how many fields of each section type there are in the +// tree under section. +func countSections(uses map[reflect.Type]int, section reflect.Type) { + for _, f := range fieldsOf(section) { + if f.section { + uses[f.typ]++ + countSections(uses, f.typ) + } + } +} + +func checkSections(t *testing.T, conf *JobSearchConfig, all map[string]any, uses map[reflect.Type]int, prefix string, section reflect.Value) { for _, f := range fieldsOf(section.Type()) { + key := joinKey(prefix, f.key) + if isSectionList(f.typ) { + checkSectionList(t, all, key, section.FieldByIndex(f.index)) + } if !f.section { continue } @@ -59,23 +83,106 @@ func checkSections(t *testing.T, all map[string]any, prefix string, section refl name := fv.Type().Name() if !fv.Addr().Type().Implements(sectionType) { t.Errorf("*%s does not implement Section", name) + continue } checkEmbedded(t, fv.Type()) - if flatten := fv.MethodByName("Flatten"); flatten.IsValid() { - keys := flatten.Call(nil)[0].Interface().(map[string]any) - for k := range keys { - if _, ok := all[joinKey(prefix, k)]; !ok { - t.Errorf("%s.Flatten() key %q does not match its cfg tag %q", - name, k, joinKey(prefix, f.key)) - } + keys := fv.Addr().Interface().(Section).Flatten() + keyPrefix := prefix + if uses[f.typ] > 1 { + keyPrefix = key + } + for k := range keys { + if _, ok := all[joinKey(keyPrefix, k)]; !ok { + t.Errorf("%s.Flatten() key %q does not match its cfg tag %q", + name, k, key) } } - checkSections(t, all, joinKey(prefix, f.key), fv) + + // The section's own keys, relative to it, with the whole + // configuration's values and kinds for them. + want := map[string]any{} + kinds := map[string]reflect.Kind{} + for k, v := range all { + if rel, ok := strings.CutPrefix(k, key+"."); ok { + want[rel] = v + kinds[rel], _ = conf.TypeOf(k) + } + } + checkSectionMethods(t, name, fv, want, kinds) + checkSections(t, conf, all, uses, key, fv) } } +// checkSectionList checks the items of the list of sections at key implement +// Section and that their methods agree with the list's value in all. +func checkSectionList(t *testing.T, all map[string]any, key string, list reflect.Value) { + name := list.Type().Elem().Name() + if !reflect.PointerTo(list.Type().Elem()).Implements(sectionType) { + t.Errorf("*%s (items of %s) does not implement Section", name, key) + return + } + items, _ := all[key].([]map[string]any) + if len(items) != list.Len() { + t.Errorf("%s has %d items; Flatten gives %d", key, list.Len(), len(items)) + return + } + for i := range list.Len() { + item, itemName := list.Index(i), fmt.Sprintf("%s[%d]", key, i) + want := flattenMap("", items[i]) + // Items are not prefixed, and their kinds are their fields' kinds. + if got := item.Addr().Interface().(Section).Flatten(); !reflect.DeepEqual(got, want) { + t.Errorf("%s.Flatten() =\n%v\nwant\n%v", itemName, got, want) + } + kinds := map[string]reflect.Kind{} + for k := range want { + _, f, _ := lookup(item, k) + kinds[k] = kindOf(f.typ) + } + checkSectionMethods(t, itemName, item, want, kinds) + } +} + +// checkSectionMethods checks section's Map, TypeOf, ValueOf and SetValue +// against want, its keys and values, and kinds, their kinds. +// Setting each key to its own value must succeed and leave section unchanged. +func checkSectionMethods(t *testing.T, name string, section reflect.Value, want map[string]any, kinds map[string]reflect.Kind) { + s := section.Addr().Interface().(Section) + if got := flattenMap("", s.Map()); !reflect.DeepEqual(got, want) { + t.Errorf("%s.Map() =\n%v\nwant\n%v", name, got, want) + } + before := section.Interface() + for k, v := range want { + if kind, err := s.TypeOf(k); err != nil || kind != kinds[k] { + t.Errorf("%s.TypeOf(%q) = %v, %v; want %v", name, k, kind, err, kinds[k]) + } + if got, err := s.ValueOf(k); err != nil || !reflect.DeepEqual(got, v) { + t.Errorf("%s.ValueOf(%q) = %#v, %v; want %#v", name, k, got, err, v) + } + if err := s.SetValue(k, v); err != nil { + t.Errorf("%s.SetValue(%q, %#v): %v", name, k, v, err) + } + } + if !reflect.DeepEqual(section.Interface(), before) { + t.Errorf("%s changed by setting its keys to their own values", name) + } +} + +// flattenMap flattens a map with nested maps, as Map gives, into dotted keys +// prefixed with prefix (if non-empty). +func flattenMap(prefix string, m map[string]any) map[string]any { + flat := map[string]any{} + for k, v := range m { + if nested, ok := v.(map[string]any); ok { + maps.Copy(flat, flattenMap(joinKey(prefix, k), nested)) + } else { + flat[joinKey(prefix, k)] = v + } + } + return flat +} + // checkEmbedded checks that structs embedded in a section, such as -// JobSearchConfigSource, implement Section too. +// JSCSourcesSource, implement Section too. func checkEmbedded(t *testing.T, section reflect.Type) { for i := 0; i < section.NumField(); i++ { sf := section.Field(i) diff --git a/internal/config/modify.go b/internal/config/modify.go index 3296f05..adf5eed 100644 --- a/internal/config/modify.go +++ b/internal/config/modify.go @@ -73,6 +73,8 @@ func parseValue(kind reflect.Kind, value string) (any, error) { return strconv.ParseUint(value, 10, 64) case reflect.Float32, reflect.Float64: return strconv.ParseFloat(value, 64) + case reflect.Slice: + return splitList(value), nil } return nil, errors.ErrUnsupported } diff --git a/internal/config/modify_test.go b/internal/config/modify_test.go index 18610a3..1c264f1 100644 --- a/internal/config/modify_test.go +++ b/internal/config/modify_test.go @@ -181,9 +181,10 @@ func TestParseValue(t *testing.T) { {reflect.Int8, "7", int64(7)}, {reflect.Uint16, "65535", uint64(65535)}, {reflect.Float32, "1.5", 1.5}, + {reflect.Slice, " a, b ,, c ", []string{"a", "b", "c"}}, } { got, err := parseValue(tt.kind, tt.value) - if err != nil || got != tt.want { + if err != nil || !reflect.DeepEqual(got, tt.want) { t.Errorf("parseValue(%v, %q) = %#v, %v; want %#v", tt.kind, tt.value, got, err, tt.want) } } @@ -202,8 +203,8 @@ func TestParseValue(t *testing.T) { } } - if _, err := parseValue(reflect.Slice, "a,b"); !errors.Is(err, errors.ErrUnsupported) { - t.Errorf("parseValue(slice) err = %v; want ErrUnsupported", err) + if _, err := parseValue(reflect.Map, "a"); !errors.Is(err, errors.ErrUnsupported) { + t.Errorf("parseValue(map) err = %v; want ErrUnsupported", err) } } diff --git a/internal/config/section.go b/internal/config/section.go index 747c7d8..366cbe6 100644 --- a/internal/config/section.go +++ b/internal/config/section.go @@ -1,12 +1,16 @@ package config import ( + "encoding/json" "errors" "fmt" "reflect" + "slices" "strconv" "strings" "time" + + "git.avsdev.uk/Go/job-search/internal/util" ) // Configuration sections are plain structs described with struct tags: @@ -22,6 +26,19 @@ import ( // // time.Duration fields report reflect.String from TypeOf and accept either a // time.Duration or a duration string such as "90s" in SetValue. +// util.ByteSize fields work the same way, with size strings such as "32KiB". +// +// []string fields report reflect.Slice and accept either a []string or a +// comma-separated string such as "a, b" in SetValue (and so in default tags). +// An empty list is stored as nil. +// +// A list of sections, []S where S is a section struct (e.g. a list of extra +// headers), is a single key reporting reflect.Slice. It is given out as a +// []map[string]any, one map per item as Map gives it, and SetValue accepts +// a []S, a list of maps (as Viper reads a YAML or JSON list of objects) or a +// JSON array string. Each item starts from its defaults; its keys are set as +// SetValue sets them, so values may be strings to parse. If the item type has +// a validator, every field of each item is validated, including ones not set. // // A section can check values before they are stored by implementing // validator on its pointer type. @@ -31,16 +48,18 @@ const ( ) var ( - errUnknownKey = errors.New("unknown configuration key") - durationType = reflect.TypeOf(time.Duration(0)) - timeType = reflect.TypeOf(time.Time{}) + errUnknownKey = errors.New("unknown configuration key") + durationType = reflect.TypeOf(time.Duration(0)) + byteSizeType = reflect.TypeOf(util.ByteSize(0)) + timeType = reflect.TypeOf(time.Time{}) + stringSliceType = reflect.TypeOf([]string{}) ) // Section is the method set of every configuration section, including // JobSearchConfig itself, implemented on the section's pointer type. Keys are // dotted paths relative to the section, e.g. "sleeper.interval" in -// JobSearchConfigSources. Durations are given out as strings such as "10m0s", -// matching TypeOf's reflect.String for them. TestSections checks every +// JSCSources. Durations and byte sizes are given out as strings such as +// "10m0s" and "32KiB", matching TypeOf's reflect.String for them. TestSections checks every // section, found by walking JobSearchConfig, implements it. type Section interface { // Map returns the section as a map of keys to values, with nested @@ -48,8 +67,9 @@ type Section interface { Map() map[string]any // Flatten returns the section as a flat map of dotted keys, each // prefixed with the section's own key (its JSCK* constant), so a - // section's keys match how its parent flattens it. JobSearchConfig and - // the embedded JobSearchConfigSource have no prefix. + // section's keys match how its parent flattens it. JobSearchConfig, the + // embedded JSCSourcesSource and types used under several keys, such as + // JSCRateLimitRule, have no prefix. Flatten() map[string]any // TypeOf returns the kind of the value at key, which is how it is read // from Viper, or an error for an unknown key. @@ -80,6 +100,13 @@ type field struct { hasDefault bool } +// isSectionList reports whether t is a list of sections, []S for a section +// struct S. +func isSectionList(t reflect.Type) bool { + return t.Kind() == reflect.Slice && t.Elem().Kind() == reflect.Struct && + t.Elem() != timeType +} + // fieldsOf lists the configuration fields of a section type in declaration // order, with the fields of embedded structs promoted into it. func fieldsOf(t reflect.Type) []field { @@ -201,19 +228,48 @@ func flattenInto(m map[string]any, prefix string, section reflect.Value) { } // exportValue returns a field's value as Map, Flatten and ValueOf give it -// out: durations become strings such as "1m30s", so values agree with kindOf -// and can be passed back to SetValue; everything else is unchanged. +// out: durations and byte sizes become strings such as "1m30s" and "32KiB", +// so values agree with kindOf and can be passed back to SetValue; []string is +// copied; a list of sections becomes a []map[string]any (nil if empty); +// everything else is unchanged. func exportValue(fv reflect.Value) any { - if fv.Type() == durationType { + switch fv.Type() { + case durationType: return time.Duration(fv.Int()).String() + case byteSizeType: + return util.ByteSize(fv.Int()).String() } - return fv.Interface() + if isSectionList(fv.Type()) { + var items []map[string]any + for i := range fv.Len() { + items = append(items, mapOfValue(fv.Index(i))) + } + return items + } + v := fv.Interface() + if s, ok := v.([]string); ok { + // A copy, so changing it doesn't change the configuration. + return slices.Clone(s) + } + return v +} + +// splitList splits a comma-separated list, trimming spaces around each item +// and dropping empty ones, so "" is an empty list. +func splitList(s string) []string { + var items []string + for item := range strings.SplitSeq(s, ",") { + if item = strings.TrimSpace(item); item != "" { + items = append(items, item) + } + } + return items } // kindOf is the kind reported by TypeOf, i.e. how the value is read from -// Viper. Durations are read as strings and parsed by SetValue. +// Viper. Durations and byte sizes are read as strings and parsed by SetValue. func kindOf(t reflect.Type) reflect.Kind { - if t == durationType { + if t == durationType || t == byteSizeType { return reflect.String } return t.Kind() @@ -265,11 +321,19 @@ func setField(section reflect.Value, f field, value any) error { } // convert turns value into a value of type t. Values already of the right -// type are used as-is; strings are parsed. +// type are used as-is, except that an empty []string becomes nil, as an empty +// comma-separated list does (see splitList); strings are parsed. Lists of +// sections are converted by convertSectionList. func convert(value any, t reflect.Type) (reflect.Value, error) { + if isSectionList(t) { + return convertSectionList(value, t) + } if value != nil { rv := reflect.ValueOf(value) if rv.Type().AssignableTo(t) { + if t == stringSliceType && rv.Len() == 0 { + return reflect.Zero(t), nil + } return rv, nil } if s, ok := value.(string); ok { @@ -279,17 +343,123 @@ func convert(value any, t reflect.Type) (reflect.Value, error) { return reflect.Value{}, fmt.Errorf("unsupported type %T", value) } +// convertSectionList turns value into a list of sections of type t: a value +// of type t, a list of maps of keys to values, or a JSON array of objects. +// nil and empty lists become a nil list. +func convertSectionList(value any, t reflect.Type) (reflect.Value, error) { + if s, ok := value.(string); ok { + var items []map[string]any + if strings.TrimSpace(s) != "" { + if err := json.Unmarshal([]byte(s), &items); err != nil { + return reflect.Value{}, fmt.Errorf("unable to convert to a list of %s: %w", t.Elem(), err) + } + } + value = items + } + if value == nil { + return reflect.Zero(t), nil + } + rv := reflect.ValueOf(value) + if rv.Kind() != reflect.Slice { + return reflect.Value{}, fmt.Errorf("unsupported type %T", value) + } + if rv.Len() == 0 { + return reflect.Zero(t), nil + } + if rv.Type().AssignableTo(t) { + for i := range rv.Len() { + if err := validateSection(rv.Index(i)); err != nil { + return reflect.Value{}, fmt.Errorf("item %d: %w", i, err) + } + } + return rv, nil + } + list := reflect.MakeSlice(t, 0, rv.Len()) + for i := range rv.Len() { + item, err := sectionFromMap(rv.Index(i).Interface(), t.Elem()) + if err != nil { + return reflect.Value{}, fmt.Errorf("item %d: %w", i, err) + } + list = reflect.Append(list, item) + } + return list, nil +} + +// sectionFromMap builds a section of type t from a map of its keys to values, +// starting from its defaults. Values that are neither strings nor of the +// field's type (e.g. numbers read from YAML) are formatted and parsed. The +// finished section is validated as a whole (see validateSection), so a key +// left out of the map can't skip its validation. +func sectionFromMap(value any, t reflect.Type) (reflect.Value, error) { + item, ok := value.(map[string]any) + if !ok { + return reflect.Value{}, fmt.Errorf("unsupported type %T", value) + } + section := reflect.New(t).Elem() + if err := applyDefaultsTo("", section); err != nil { + return reflect.Value{}, err + } + for k, v := range item { + s, f, err := lookup(section, k) + if err != nil { + return reflect.Value{}, fmt.Errorf("%w %q", err, k) + } + if _, isString := v.(string); v != nil && !isString && + !reflect.TypeOf(v).AssignableTo(f.typ) { + v = fmt.Sprint(v) + } + if err := setField(s, f, v); err != nil { + return reflect.Value{}, fmt.Errorf("%s: %w", k, err) + } + } + if err := validateSection(section); err != nil { + return reflect.Value{}, err + } + return section, nil +} + +// validateSection runs the section's validator, if it has one, on every value +// field directly in it, so values that were never set (defaults or zero +// values) are checked too. +func validateSection(section reflect.Value) error { + v, ok := reflect.New(section.Type()).Interface().(validator) + if !ok { + return nil + } + for _, f := range fieldsOf(section.Type()) { + if f.section { + continue + } + if err := v.validate(f.key, section.FieldByIndex(f.index).Interface()); err != nil { + return fmt.Errorf("%s: %w", f.key, err) + } + } + return nil +} + // parseString parses s as a value of type t. It supports strings, bools, -// signed and unsigned integers, floats and time.Duration. +// signed and unsigned integers, floats, time.Duration, ByteSize and []string +// (from a comma-separated list, see splitList). func parseString(s string, t reflect.Type) (reflect.Value, error) { v := reflect.New(t).Elem() - if t == durationType { + switch t { + case durationType: d, err := time.ParseDuration(s) if err != nil { return reflect.Value{}, fmt.Errorf("unable to convert to duration: %w", err) } v.SetInt(int64(d)) return v, nil + case byteSizeType: + b, err := util.ParseByteSize(s) + if err != nil { + return reflect.Value{}, fmt.Errorf("unable to convert to byte size: %w", err) + } + v.SetInt(int64(b)) + return v, nil + case stringSliceType: + v.Set(reflect.ValueOf(splitList(s))) + return v, nil } var err error diff --git a/internal/config/section_test.go b/internal/config/section_test.go index bb16bef..3dd6443 100644 --- a/internal/config/section_test.go +++ b/internal/config/section_test.go @@ -6,6 +6,8 @@ import ( "strings" "testing" "time" + + "git.avsdev.uk/Go/job-search/internal/util" ) // Test sections exercising every field shape section.go supports. @@ -13,6 +15,13 @@ type testInner struct { Name string `cfg:"name" default:"inner"` } +// testItem is the item type of a list of sections. +type testItem struct { + Name string `cfg:"name" default:"item"` + N int `cfg:"n"` + Wait time.Duration `cfg:"wait"` +} + type testEmbedded struct { On bool `cfg:"on" default:"true"` } @@ -27,8 +36,10 @@ type testSection struct { Float float64 `cfg:"float" default:"1.5"` Float32 float32 `cfg:"float32"` Dur time.Duration `cfg:"dur" default:"90s"` + Size util.ByteSize `cfg:"size" default:"32KiB"` When time.Time `cfg:"when"` Slice []string `cfg:"slice"` + Items []testItem `cfg:"items"` Inner testInner `cfg:"inner"` NoTag string Skipped string `cfg:"-"` @@ -52,7 +63,7 @@ func (c *testValidated) validate(key string, value any) error { func TestKeysOf(t *testing.T) { got := keysOf(&testSection{}) want := []string{"on", "str", "int", "int8", "uint", "uint16", "float", - "float32", "dur", "when", "slice", "inner"} + "float32", "dur", "size", "when", "slice", "items", "inner"} if !reflect.DeepEqual(got, want) { t.Errorf("keysOf = %v; want %v", got, want) } @@ -70,6 +81,9 @@ func TestFieldsOfSkipsUntagged(t *testing.T) { if f.key == "inner" && !f.section { t.Error("struct field not treated as a section") } + if f.key == "items" && f.section { + t.Error("list of sections treated as a section") + } } } @@ -137,6 +151,7 @@ func TestApplyDefaults(t *testing.T) { Uint: 7, Float: 1.5, Dur: 90 * time.Second, + Size: 32 << 10, Inner: testInner{Name: "inner"}, } if !reflect.DeepEqual(s, want) { @@ -174,9 +189,11 @@ func TestTypeOf(t *testing.T) { {"uint", reflect.Uint}, {"float", reflect.Float64}, {"on", reflect.Bool}, - {"dur", reflect.String}, // durations are read as strings + {"dur", reflect.String}, // durations are read as strings + {"size", reflect.String}, // byte sizes are read as strings {"when", reflect.Struct}, {"slice", reflect.Slice}, + {"items", reflect.Slice}, {"inner.name", reflect.String}, } for _, tt := range tests { @@ -227,12 +244,38 @@ func TestSetValue(t *testing.T) { {"float", "2.25", 2.25}, {"float32", "0.5", float32(0.5)}, {"dur", "1h30m", "1h30m0s"}, // durations are given out as strings + {"size", "1 mb", "1MB"}, // byte sizes are given out as strings + {"size", "1.5KiB", "1536B"}, + {"size", "4096", "4KiB"}, + {"slice", " a, b ,, c ", []string{"a", "b", "c"}}, + {"slice", "", []string(nil)}, {"inner.name", "nested", "nested"}, // values of the field's type are used as-is {"int", 5, 5}, {"on", true, true}, {"dur", 3 * time.Second, "3s"}, + {"size", util.ByteSize(2048), "2KiB"}, {"slice", []string{"a", "b"}, []string{"a", "b"}}, + {"slice", []string{}, []string(nil)}, // empty lists are stored as nil + // lists of sections are given out as lists of maps; items start + // from their defaults and their values are parsed + {"items", `[{"n": "1", "wait": "2s"}, {"name": "b", "n": 2}]`, []map[string]any{ + {"name": "item", "n": 1, "wait": "2s"}, + {"name": "b", "n": 2, "wait": "0s"}, + }}, + {"items", []any{map[string]any{"name": "a", "n": 3}}, []map[string]any{ + {"name": "a", "n": 3, "wait": "0s"}, + }}, + {"items", []map[string]any{{"name": "a"}}, []map[string]any{ + {"name": "a", "n": 0, "wait": "0s"}, + }}, + {"items", []testItem{{Name: "a", N: 4}}, []map[string]any{ + {"name": "a", "n": 4, "wait": "0s"}, + }}, + {"items", "", []map[string]any(nil)}, + {"items", "[]", []map[string]any(nil)}, + {"items", []testItem{}, []map[string]any(nil)}, + {"items", nil, []map[string]any(nil)}, {"when", time.Unix(0, 0).UTC(), time.Unix(0, 0).UTC()}, } for _, tt := range tests { @@ -261,7 +304,14 @@ func TestSetValueErrors(t *testing.T) { {"float", "x", "unable to convert"}, {"on", "maybe", "unable to convert"}, {"dur", "soon", "unable to convert to duration"}, - {"slice", "a,b", "unsupported type string"}, + {"items", "not json", "unable to convert to a list of config.testItem"}, + {"items", `[{"bogus": 1}]`, `item 0: unknown configuration key "bogus"`}, + {"items", `[{"n": "x"}]`, "item 0: n: unable to convert"}, + {"items", []any{"a"}, "item 0: unsupported type string"}, + {"items", 5, "unsupported type int"}, + {"size", "lots", "unable to convert to byte size"}, + {"size", "-1KB", "unable to convert to byte size"}, + {"size", "16EiB", "too large"}, {"when", "2024-01-01", "unsupported type string"}, {"int", 1.5, "unsupported type float64"}, {"int", int64(1), "unsupported type int64"}, @@ -281,6 +331,20 @@ func TestSetValueErrors(t *testing.T) { } } +// A []string is given out as a copy, so changing it leaves the section as it +// was. +func TestValueOfSliceIsACopy(t *testing.T) { + s := testSection{Slice: []string{"a", "b"}} + got, err := valueOf(s, "slice") + if err != nil { + t.Fatal(err) + } + got.([]string)[0] = "changed" + if s.Slice[0] != "a" { + t.Errorf("changing valueOf's slice changed the section: %v", s.Slice) + } +} + func TestSetValueValidator(t *testing.T) { var s testValidated if err := setValue(&s, "value", "good"); err != nil || s.Value != "good" { diff --git a/internal/httpx/httpx_test.go b/internal/httpx/httpx_test.go new file mode 100644 index 0000000..bbda3ac --- /dev/null +++ b/internal/httpx/httpx_test.go @@ -0,0 +1,84 @@ +package httpx + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" +) + +func TestWriteProblem(t *testing.T) { + req := httptest.NewRequest("GET", "/api/items/1", nil) + req = req.WithContext(WithRequestID(req.Context(), "req-123")) + rec := httptest.NewRecorder() + WriteProblem(rec, req, http.StatusNotFound, "no such item") + + if rec.Code != http.StatusNotFound { + t.Errorf("status = %d; want 404", rec.Code) + } + if ct := rec.Header().Get("Content-Type"); ct != "application/problem+json" { + t.Errorf("Content-Type = %q", ct) + } + if nosniff := rec.Header().Get("X-Content-Type-Options"); nosniff != "nosniff" { + t.Errorf("X-Content-Type-Options = %q", nosniff) + } + var p Problem + if err := json.Unmarshal(rec.Body.Bytes(), &p); err != nil { + t.Fatalf("body %q: %v", rec.Body, err) + } + want := Problem{ + Type: "about:blank", + Status: http.StatusNotFound, + Title: "Not Found", + Detail: "no such item", + RequestID: "req-123", + } + if p != want { + t.Errorf("problem = %+v; want %+v", p, want) + } +} + +// Without a request (or a request ID) the problem has no requestId, and an +// empty detail is left out. +func TestWriteProblemMinimal(t *testing.T) { + for _, req := range []*http.Request{nil, httptest.NewRequest("GET", "/", nil)} { + rec := httptest.NewRecorder() + WriteProblem(rec, req, http.StatusTooManyRequests, "") + var body map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil { + t.Fatalf("body %q: %v", rec.Body, err) + } + want := map[string]any{ + "type": "about:blank", + "status": float64(http.StatusTooManyRequests), + "title": "Too Many Requests", + } + if len(body) != len(want) { + t.Errorf("request %v: body = %v; want %v", req != nil, body, want) + } + for k, v := range want { + if body[k] != v { + t.Errorf("request %v: %s = %v; want %v", req != nil, k, body[k], v) + } + } + } +} + +func TestRequestID(t *testing.T) { + ctx := context.Background() + if id := RequestID(ctx); id != "" { + t.Errorf("RequestID(empty) = %q; want empty", id) + } + ctx = WithRequestID(ctx, "abc") + if id := RequestID(ctx); id != "abc" { + t.Errorf("RequestID = %q; want abc", id) + } + // The latest ID wins; a value of another type under the key is ignored. + if id := RequestID(WithRequestID(ctx, "def")); id != "def" { + t.Errorf("RequestID after replacing = %q; want def", id) + } + if id := RequestID(context.WithValue(ctx, requestIDKey, 42)); id != "" { + t.Errorf("RequestID(int) = %q; want empty", id) + } +} diff --git a/internal/httpx/problem.go b/internal/httpx/problem.go new file mode 100644 index 0000000..abbe97e --- /dev/null +++ b/internal/httpx/problem.go @@ -0,0 +1,34 @@ +// Package httpx holds small HTTP helpers shared by middleware and handlers. +package httpx + +import ( + "encoding/json" + "net/http" +) + +// Problem is an RFC 9457 problem details body. +type Problem struct { + Type string `json:"type,omitempty"` + Status int `json:"status"` + Title string `json:"title"` + Detail string `json:"detail,omitempty"` + Instance string `json:"instance,omitempty"` + RequestID string `json:"requestId,omitempty"` +} + +// WriteProblem writes a problem+json error response. +func WriteProblem(w http.ResponseWriter, r *http.Request, status int, detail string) { + p := Problem{ + Type: "about:blank", + Title: http.StatusText(status), + Status: status, + Detail: detail, + } + if r != nil { + p.RequestID = RequestID(r.Context()) + } + w.Header().Set("Content-Type", "application/problem+json") + w.Header().Set("X-Content-Type-Options", "nosniff") + w.WriteHeader(status) + json.NewEncoder(w).Encode(p) //nolint:errcheck,gosec // the client may have gone; nothing to do +} diff --git a/internal/httpx/request_id.go b/internal/httpx/request_id.go new file mode 100644 index 0000000..5c61845 --- /dev/null +++ b/internal/httpx/request_id.go @@ -0,0 +1,20 @@ +package httpx + +import ( + "context" +) + +type ctxKey int + +const requestIDKey ctxKey = iota + +// WithRequestID stores a request ID in the context. +func WithRequestID(ctx context.Context, id string) context.Context { + return context.WithValue(ctx, requestIDKey, id) +} + +// RequestID returns the request ID stored in ctx, or "". +func RequestID(ctx context.Context) string { + id, _ := ctx.Value(requestIDKey).(string) + return id +} diff --git a/internal/middleware/ratelimit.go b/internal/middleware/ratelimit.go new file mode 100644 index 0000000..dd8ad80 --- /dev/null +++ b/internal/middleware/ratelimit.go @@ -0,0 +1,144 @@ +package middleware + +import ( + "context" + "hash/maphash" + "math" + "net/http" + "strconv" + "sync" + "time" + + "golang.org/x/time/rate" + + "git.avsdev.uk/Go/job-search/internal/httpx" +) + +// Limiter is the rate limiting hook. Implementations decide whether a request +// identified by key may proceed. Swap in a Redis/memcached-backed limiter for +// multi-instance deployments by implementing this interface. +type Limiter interface { + Allow(key string) (allowed bool, retryAfter time.Duration) +} + +// LimiterFunc adapts a function to the Limiter interface. +type LimiterFunc func(key string) (bool, time.Duration) + +func (f LimiterFunc) Allow(key string) (bool, time.Duration) { return f(key) } + +// KeyFunc derives the rate limit key from a request. +type KeyFunc func(*http.Request) string + +// KeyByIP keys on the real client IP (after proxy processing). +func KeyByIP(r *http.Request) string { return ClientIP(r) } + +// RateLimit turns a Limiter into middleware. Rejected requests get 429 with +// a Retry-After header. A nil limiter yields a pass-through middleware. +func RateLimit(l Limiter, key KeyFunc) func(http.Handler) http.Handler { + if l == nil { + return func(next http.Handler) http.Handler { return next } + } + if key == nil { + key = KeyByIP + } + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ok, retry := l.Allow(key(r)) + if !ok { + secs := int(math.Ceil(retry.Seconds())) + if secs < 1 { + secs = 1 + } + w.Header().Set("Retry-After", strconv.Itoa(secs)) + httpx.WriteProblem(w, r, http.StatusTooManyRequests, "rate limit exceeded") + return + } + next.ServeHTTP(w, r) + }) + } +} + +// KeyedLimiter is an in-memory token bucket per key, sharded to reduce lock +// contention, with idle-entry eviction. +type KeyedLimiter struct { + limit rate.Limit + burst int + ttl time.Duration + seed maphash.Seed + shards [32]limiterShard +} + +type limiterShard struct { + mu sync.Mutex + entries map[string]*limiterEntry +} + +type limiterEntry struct { + lim *rate.Limiter + last time.Time +} + +// NewKeyedLimiter creates a limiter allowing rps sustained requests per +// second with the given burst. A janitor goroutine evicts entries idle for +// longer than ttl until ctx is cancelled. +func NewKeyedLimiter(ctx context.Context, rps float64, burst int, ttl time.Duration) *KeyedLimiter { + k := &KeyedLimiter{limit: rate.Limit(rps), burst: burst, ttl: ttl, seed: maphash.MakeSeed()} + for i := range k.shards { + k.shards[i].entries = make(map[string]*limiterEntry) + } + if ttl > 0 { + go k.janitor(ctx) + } + return k +} + +func (k *KeyedLimiter) shard(key string) *limiterShard { + return &k.shards[maphash.String(k.seed, key)%uint64(len(k.shards))] +} + +// Allow implements Limiter. +func (k *KeyedLimiter) Allow(key string) (bool, time.Duration) { + now := time.Now() + s := k.shard(key) + s.mu.Lock() + e := s.entries[key] + if e == nil { + e = &limiterEntry{lim: rate.NewLimiter(k.limit, k.burst)} + s.entries[key] = e + } + e.last = now + lim := e.lim + s.mu.Unlock() + + res := lim.ReserveN(now, 1) + if !res.OK() { + return false, time.Second + } + if d := res.DelayFrom(now); d > 0 { + res.CancelAt(now) + return false, d + } + return true, 0 +} + +func (k *KeyedLimiter) janitor(ctx context.Context) { + t := time.NewTicker(k.ttl / 2) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case now := <-t.C: + for i := range k.shards { + s := &k.shards[i] + s.mu.Lock() + for key, e := range s.entries { + if now.Sub(e.last) > k.ttl { + delete(s.entries, key) + } + } + s.mu.Unlock() + } + } + } +} diff --git a/internal/middleware/ratelimit_test.go b/internal/middleware/ratelimit_test.go new file mode 100644 index 0000000..06f127f --- /dev/null +++ b/internal/middleware/ratelimit_test.go @@ -0,0 +1,24 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestRateLimit(t *testing.T) { + l := NewKeyedLimiter(t.Context(), 1, 2, 0) + h := RateLimit(l, nil)(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + var last int + for i := 0; i < 4; i++ { + rec := httptest.NewRecorder() + h.ServeHTTP(rec, httptest.NewRequest("GET", "/", nil)) + last = rec.Code + if last == http.StatusTooManyRequests && rec.Header().Get("Retry-After") == "" { + t.Fatal("missing Retry-After") + } + } + if last != http.StatusTooManyRequests { + t.Fatalf("want 429, got %d", last) + } +} diff --git a/internal/middleware/real_ip.go b/internal/middleware/real_ip.go new file mode 100644 index 0000000..55430fe --- /dev/null +++ b/internal/middleware/real_ip.go @@ -0,0 +1,236 @@ +package middleware + +import ( + "context" + "net/http" + "net/netip" + "strings" +) + +// ========== RealIP - ClientInfo ========== + +// ClientInfo describes the client as seen after proxy processing. +type ClientInfo struct { + IP netip.Addr // best-effort real client IP + Peer netip.Addr // immediate TCP peer (after PROXY protocol, if used) + Scheme string // "http" or "https", honouring trusted X-Forwarded-Proto + ViaProxy bool // true when IP came from a trusted forwarding header +} + +// Client returns the ClientInfo stored by RealIP. Without RealIP it falls back +// to the connection itself (see peerInfo) and ignores forwarding headers. +func Client(r *http.Request) ClientInfo { + if ci, ok := r.Context().Value(clientInfoKey).(ClientInfo); ok { + return ci + } + return peerInfo(r) +} + +// ClientIP returns the client IP as a string ("" if unknown). +func ClientIP(r *http.Request) string { + ip := Client(r).IP + if !ip.IsValid() { + return "" + } + return ip.String() +} + +// IsSecure reports whether the request reached the edge over HTTPS. Without +// RealIP it only reports whether the connection to this server is TLS. +func IsSecure(r *http.Request) bool { + return Client(r).Scheme == "https" +} + +// peerInfo builds a ClientInfo from the connection alone: the TCP peer (after +// PROXY protocol, if used) and whether this server terminated TLS. +func peerInfo(r *http.Request) ClientInfo { + peer := addrFromRemote(r.RemoteAddr) + + scheme := "http" + if r.TLS != nil { + scheme = "https" + } + + return ClientInfo{ + IP: peer, + Peer: peer, + Scheme: scheme, + } +} + +// ========== RealIP ========== + +// RealIPOptions configures RealIP. +type RealIPOptions struct { + Trusted []netip.Prefix + // Headers is processed in increasing priority order - the last header in the + // list that is found is the header used to determine a client IP address. + Headers []string // e.g. X-Forwarded-For, Forwarded, X-Real-IP + StripUntrusted bool +} + +type ctxKey int + +const clientInfoKey ctxKey = iota + +var forwardingHeaders = []string{ + "X-Forwarded-For", + "X-Forwarded-Proto", + "X-Forwarded-Host", + "X-Forwarded-Port", + "X-Forwarded-Ssl", + "Forwarded", + "X-Real-IP", +} + +// RealIP determines the client IP and scheme. Forwarding headers are only +// honoured when the TCP peer is a trusted proxy; the X-Forwarded-For and +// Forwarded chains are walked right-to-left skipping trusted hops, so a +// client cannot spoof its address by prepending entries. +// +// PROXY protocol is handled at the listener (see server package), so by the +// time this runs r.RemoteAddr already holds the address from the PROXY header. +func RealIP(opts RealIPOptions) func(http.Handler) http.Handler { + trusted := func(a netip.Addr) bool { + for _, p := range opts.Trusted { + if p.Contains(a) { + return true + } + } + return false + } + + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + info := peerInfo(r) + + if info.Peer.IsValid() && trusted(info.Peer) { + for _, h := range opts.Headers { + if ip, ok := clientFromHeader(r.Header, h, trusted); ok { + info.IP = ip + info.ViaProxy = true + break + } + } + if p := forwardedProto(r.Header); p == "http" || p == "https" { + info.Scheme = p + } + } else if opts.StripUntrusted { + for _, h := range forwardingHeaders { + r.Header.Del(h) + } + } + + ctx := context.WithValue(r.Context(), clientInfoKey, info) + next.ServeHTTP(w, r.WithContext(ctx)) + }) + } +} + +func clientFromHeader( + h http.Header, name string, trusted func(netip.Addr) bool, +) (netip.Addr, bool) { + switch http.CanonicalHeaderKey(name) { + case "X-Forwarded-For": + return walkChain(splitList(h.Values("X-Forwarded-For")), trusted) + + case "Forwarded": + var chain []string + for _, el := range splitList(h.Values("Forwarded")) { + if v, ok := forwardedParam(el, "for"); ok { + chain = append(chain, v) + } else { + chain = append(chain, "") // element without for= breaks the chain + } + } + return walkChain(chain, trusted) + + case "X-Real-Ip": + if v := strings.TrimSpace(h.Get("X-Real-IP")); v != "" { + return parseIPToken(v) + } + } + + return netip.Addr{}, false +} + +// walkChain returns the right-most address that is not a trusted proxy. If +// every hop is trusted, the left-most is returned. Any unparsable hop +// ("unknown", obfuscated identifiers, garbage) stops the walk. +func walkChain( + chain []string, trusted func(netip.Addr) bool, +) (netip.Addr, bool) { + var last netip.Addr + for i := len(chain) - 1; i >= 0; i-- { + ip, ok := parseIPToken(chain[i]) + if !ok { + return netip.Addr{}, false + } + if !trusted(ip) { + return ip, true + } + last = ip + } + return last, last.IsValid() +} + +func forwardedProto(h http.Header) string { + if v := h.Get("X-Forwarded-Proto"); v != "" { + first, _, _ := strings.Cut(v, ",") + return strings.ToLower(strings.TrimSpace(first)) + } + + if els := splitList(h.Values("Forwarded")); len(els) > 0 { + if v, ok := forwardedParam(els[0], "proto"); ok { + return strings.ToLower(v) + } + } + + return "" +} + +// forwardedParam extracts a parameter from one Forwarded HTTP Extension +// element (RFC 7239). +func forwardedParam(element, key string) (string, bool) { + for _, pair := range strings.Split(element, ";") { + k, v, ok := strings.Cut(strings.TrimSpace(pair), "=") + if ok && strings.EqualFold(k, key) { + return strings.Trim(strings.TrimSpace(v), `"`), true + } + } + return "", false +} + +func splitList(values []string) []string { + var out []string + for _, v := range values { + for _, part := range strings.Split(v, ",") { + if part = strings.TrimSpace(part); part != "" { + out = append(out, part) + } + } + } + return out +} + +// parseIPToken accepts "1.2.3.4", "1.2.3.4:80", "::1", "[::1]" and "[::1]:80". +func parseIPToken(s string) (netip.Addr, bool) { + s = strings.Trim(strings.TrimSpace(s), `"`) + if ap, err := netip.ParseAddrPort(s); err == nil { + return ap.Addr().Unmap(), true + } + + s = strings.TrimSuffix(strings.TrimPrefix(s, "["), "]") + if a, err := netip.ParseAddr(s); err == nil { + return a.Unmap().WithZone(""), true + } + + return netip.Addr{}, false +} + +func addrFromRemote(remote string) netip.Addr { + if a, ok := parseIPToken(remote); ok { + return a + } + return netip.Addr{} +} diff --git a/internal/middleware/real_ip_test.go b/internal/middleware/real_ip_test.go new file mode 100644 index 0000000..896751c --- /dev/null +++ b/internal/middleware/real_ip_test.go @@ -0,0 +1,56 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "net/netip" + "testing" +) + +// serveRealIP runs a request from peer with the given headers through RealIP +// and returns the ClientInfo the next handler sees. +func serveRealIP(t *testing.T, opts RealIPOptions, peer string, h map[string]string) ClientInfo { + t.Helper() + var got ClientInfo + next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { got = Client(r) }) + req := httptest.NewRequest("GET", "/", nil) + req.RemoteAddr = peer + for k, v := range h { + req.Header.Set(k, v) + } + RealIP(opts)(next).ServeHTTP(httptest.NewRecorder(), req) + return got +} + +var defaultHeaders = []string{"X-Forwarded-For", "Forwarded", "X-Real-IP"} + +func TestRealIP(t *testing.T) { + opts := RealIPOptions{ + Trusted: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}, + Headers: defaultHeaders, + } + // Spoofed left-most entry is ignored; right-most untrusted wins. + c := serveRealIP(t, opts, "127.0.0.1:1234", map[string]string{ + "X-Forwarded-For": "6.6.6.6, 203.0.113.9, 127.0.0.1", "X-Forwarded-Proto": "https", + }) + if c.IP.String() != "203.0.113.9" || c.Scheme != "https" || !c.ViaProxy { + t.Fatalf("xff: %+v", c) + } + c = serveRealIP(t, opts, "127.0.0.1:1234", map[string]string{ + "Forwarded": `for="[2001:db8::7]:4711";proto=https`, + }) + if c.IP.String() != "2001:db8::7" { + t.Fatalf("forwarded: %+v", c) + } +} + +func TestUntrustedPeerHeadersIgnored(t *testing.T) { + opts := RealIPOptions{ + Trusted: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}, + Headers: defaultHeaders, + } + c := serveRealIP(t, opts, "127.0.0.1:1234", map[string]string{"X-Forwarded-For": "203.0.113.9"}) + if c.IP.String() != "127.0.0.1" || c.ViaProxy { + t.Fatalf("untrusted XFF honoured: %+v", c) + } +} diff --git a/internal/util/bytesize.go b/internal/util/bytesize.go new file mode 100644 index 0000000..9403af8 --- /dev/null +++ b/internal/util/bytesize.go @@ -0,0 +1,53 @@ +package util + +import ( + "fmt" + "math" + "strconv" + + "github.com/dustin/go-humanize" +) + +// ByteSize is a size in bytes, configured as a string such as "32KiB", "1MB" +// or "1.5 GB" (see ParseByteSize). KB, MB, ... are powers of 1000 and KiB, +// MiB, ... powers of 1024; a bare number is bytes. +type ByteSize int64 + +// byteUnits are the units String may use, largest first. +var byteUnits = []struct { + name string + size int64 +}{ + {"EiB", humanize.EiByte}, {"EB", humanize.EByte}, + {"PiB", humanize.PiByte}, {"PB", humanize.PByte}, + {"TiB", humanize.TiByte}, {"TB", humanize.TByte}, + {"GiB", humanize.GiByte}, {"GB", humanize.GByte}, + {"MiB", humanize.MiByte}, {"MB", humanize.MByte}, + {"KiB", humanize.KiByte}, {"KB", humanize.KByte}, +} + +// ParseByteSize parses a size such as "32KiB", "1 MB" or "4096". Unit names +// are case-insensitive. +func ParseByteSize(s string) (ByteSize, error) { + n, err := humanize.ParseBytes(s) + if err != nil { + return 0, err + } + if n > math.MaxInt64 { + return 0, fmt.Errorf("byte size %q is too large", s) + } + return ByteSize(n), nil +} + +// String formats b exactly, in the largest unit that divides it (e.g. "32KiB", +// "1MB"), or in bytes ("1500B"), so it parses back to the same value. +func (b ByteSize) String() string { + if b != 0 { + for _, u := range byteUnits { + if int64(b)%u.size == 0 { + return strconv.FormatInt(int64(b)/u.size, 10) + u.name + } + } + } + return strconv.FormatInt(int64(b), 10) + "B" +} diff --git a/internal/util/bytesize_test.go b/internal/util/bytesize_test.go new file mode 100644 index 0000000..3b118c6 --- /dev/null +++ b/internal/util/bytesize_test.go @@ -0,0 +1,27 @@ +package util + +import "testing" + +func TestByteSizeString(t *testing.T) { + tests := []struct { + b ByteSize + want string + }{ + {0, "0B"}, + {1500, "1500B"}, + {1024, "1KiB"}, + {1000, "1KB"}, + {32768, "32KiB"}, + {1000000, "1MB"}, + {1024000, "1000KiB"}, + {3 << 30, "3GiB"}, + } + for _, tt := range tests { + if got := tt.b.String(); got != tt.want { + t.Errorf("ByteSize(%d).String() = %q; want %q", int64(tt.b), got, tt.want) + } + if back, err := ParseByteSize(tt.want); err != nil || back != tt.b { + t.Errorf("ParseByteSize(%q) = %d, %v; want %d", tt.want, back, err, tt.b) + } + } +} diff --git a/internal/websocket/client.go b/internal/websocket/client.go new file mode 100644 index 0000000..f5b9519 --- /dev/null +++ b/internal/websocket/client.go @@ -0,0 +1,300 @@ +package websocket + +import ( + "context" + "crypto/rand" + "encoding/json" + "errors" + "log/slog" + "sync" + "sync/atomic" + "time" + + "github.com/coder/websocket" + "golang.org/x/time/rate" +) + +// clientSendBuffer is how many messages are queued for a client before it +// counts as a slow consumer and is closed. +const clientSendBuffer = 256 + +// logMessageLimit is how much of a message is logged; longer messages are +// cut short and logged as truncated. +const logMessageLimit = 1 << 10 + +// msgConn is the connection a client reads and writes messages on: a +// *websocket.Conn, or a tcpConn for plain TCP clients (see tcp.go). +type msgConn interface { + Read(ctx context.Context) (websocket.MessageType, []byte, error) + Write(ctx context.Context, messageType websocket.MessageType, payload []byte) error + Ping(ctx context.Context) error + Close(code websocket.StatusCode, reason string) error + CloseNow() error +} + +// client is one connection registered with the hub. Its read loop handles +// incoming envelopes, its write loop sends the queued messages and its ping +// loop checks an idle connection is still alive. +type client struct { + hub *Hub + conn msgConn + send chan []byte // queued messages, sent by writeLoop + id string + ip string + lastSeen atomic.Int64 // unix nanoseconds; set by the read and ping loops + cancel context.CancelFunc + limiter *rate.Limiter // limits the messages the client sends + closeOnce sync.Once +} + +// errUnsupportedData is returned by readLoop (and a tcpConn's Write) for a +// binary message. +var errUnsupportedData = errors.New("binary frames are not supported") + +// newClient creates a client for conn. cancel cancels the context the client +// runs with; the client calls it when it finishes or is closed. The client's +// messages are limited by websocket.message_rate_limit if it is enabled. +func newClient( + h *Hub, conn msgConn, ip string, cancel context.CancelFunc, +) *client { + limiter := rate.NewLimiter(rate.Inf, 0) + if rule := h.conf.MessageRateLimit; rule.Enabled && rule.PerSecond > 0 { + limiter = rate.NewLimiter( + rate.Limit(rule.PerSecond), + max(rule.Burst, 1), + ) + } + + c := &client{ + hub: h, + conn: conn, + id: rand.Text()[:12], + ip: ip, + send: make(chan []byte, clientSendBuffer), + cancel: cancel, + limiter: limiter, + } + c.lastSeen.Store(time.Now().UnixNano()) + + return c +} + +// run starts the writer and keepalive goroutines and reads until the +// connection ends. It returns once the connection is finished. +func (c *client) run(ctx context.Context) { + var wg sync.WaitGroup + wg.Add(2) + go func() { defer wg.Done(); c.writeLoop(ctx) }() + go func() { defer wg.Done(); c.pingLoop(ctx) }() + + c.sendEnvelope(Envelope{ + Type: TypeWelcome, + Payload: mustRaw(map[string]any{ + "clientId": c.id, + }), + }) + + err := c.readLoop(ctx) + switch { + case err == nil: + case websocket.CloseStatus(err) != -1: + // The peer closed the connection, or close() closed it. + case errors.Is(err, errUnsupportedData): + c.close(websocket.StatusUnsupportedData, "text frames only") + case errors.Is(err, context.Canceled): + // The client was cancelled: by close() (shutdown or a slow + // consumer), after a failed write or ping, or by the caller. + c.conn.CloseNow() //nolint:errcheck,gosec // finished with; it may already be closed + default: + c.hub.log.Debug("read error", "client", c.id, "error", err) + c.conn.CloseNow() //nolint:errcheck,gosec // the connection has failed anyway + } + + c.cancel() + wg.Wait() +} + +// readLoop reads messages until the connection ends, and returns why. +// Messages over the rate limit or that aren't valid JSON envelopes get an +// error reply; the rest are passed to handle. +func (c *client) readLoop(ctx context.Context) error { + for { + messageType, data, err := c.conn.Read(ctx) + if err != nil { + return err + } + c.logMessage(ctx, "message received", data) + c.lastSeen.Store(time.Now().UnixNano()) + if messageType != websocket.MessageText { + return errUnsupportedData + } + // Check if the client is rate limited for early discard + if !c.limiter.Allow() { + c.sendError("", "rate limited") + continue + } + // Unpack the envelope and pass to handle to process if valid + var env Envelope + if err := json.Unmarshal(data, &env); err != nil { + c.sendError("", "invalid JSON envelope") + continue + } + c.handle(env) + } +} + +// handle answers one envelope from the client: ping and echo, or an error +// for any other type. +func (c *client) handle(env Envelope) { + switch env.Type { + case TypePing: + c.sendEnvelope(Envelope{ + Type: TypePong, + ID: env.ID, + Payload: mustRaw(map[string]any{"time": time.Now().UTC()}), + }) + + case TypeEcho: + c.sendEnvelope(Envelope{ + Type: TypeEcho, + ID: env.ID, + Payload: env.Payload, + }) + + default: + c.sendError(env.ID, "unknown message type") + } +} + +// writeLoop sends queued messages, each within the write timeout (no limit if +// it is 0 or less), until ctx is cancelled. A failed write cancels the client. +func (c *client) writeLoop(ctx context.Context) { + for { + select { + case <-ctx.Done(): + return + case msg := <-c.send: + writeCtx, cancel := withTimeout(ctx, c.hub.conf.WriteTimeout) + err := c.conn.Write(writeCtx, websocket.MessageText, msg) + cancel() + if err != nil { + c.hub.log.Debug("write failed", "client", c.id, "error", err) + c.cancel() + return + } + c.logMessage(ctx, "message sent", msg) + } + } +} + +// pingLoop pings the connection when it has been idle for the ping interval; +// no pong within the pong timeout (no limit if it is 0 or less) drops the +// connection. Ping needs the concurrent reader in readLoop to see pongs. +func (c *client) pingLoop(ctx context.Context) { + pingInterval := c.hub.conf.PingInterval + pongTimeout := c.hub.conf.PongTimeout + + if pingInterval <= 0 { + return + } + + ticker := time.NewTicker(pingInterval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + // Don't ping if the lastSeen is more recent than the ping interval + lastSeen := time.Unix(0, c.lastSeen.Load()) + if time.Now().Before(lastSeen.Add(pingInterval)) { + continue + } + pingCtx, cancel := withTimeout(ctx, pongTimeout) + err := c.conn.Ping(pingCtx) + cancel() + if err != nil { + c.hub.log.Debug("ping failed", "client", c.id, "error", err) + c.conn.CloseNow() //nolint:errcheck,gosec // dropping the client anyway + c.cancel() + return + } + c.lastSeen.Store(time.Now().UnixNano()) + } + } +} + +// trySend queues msg without blocking; false means the buffer is full. +func (c *client) trySend(msg []byte) bool { + select { + case c.send <- msg: + return true + default: + return false + } +} + +// sendEnvelope queues env for the client, closing the client as a slow +// consumer if its queue is full. +func (c *client) sendEnvelope(env Envelope) { + encoded, err := json.Marshal(env) + if err != nil { + c.hub.log.Error("marshal", "client", c.id, "error", err) + return + } + if !c.trySend(encoded) { + c.hub.log.Warn("dropping slow client", "client", c.id, "ip", c.ip) + go c.close(websocket.StatusPolicyViolation, "slow consumer") + } +} + +// sendError queues an error envelope replying to the message with id (none +// if empty). +func (c *client) sendError(id, msg string) { + c.sendEnvelope(Envelope{ + Type: TypeError, + ID: id, + Error: msg, + }) +} + +// close performs a graceful close handshake once, then cancels the client. +// It may block for the handshake, so call it from its own goroutine unless +// the read loop has already exited. +func (c *client) close(code websocket.StatusCode, reason string) { + c.closeOnce.Do(func() { + c.conn.Close(code, reason) //nolint:errcheck,gosec // the client is cancelled anyway + c.cancel() + }) +} + +// logMessage logs a message sent to or received from the client at debug +// level, with its size; at most logMessageLimit bytes of it are logged. It +// does nothing unless debug logging is enabled. +func (c *client) logMessage(ctx context.Context, event string, data []byte) { + if !c.hub.log.Enabled(ctx, slog.LevelDebug) { + return + } + logged, truncated := data, false + if len(logged) > logMessageLimit { + logged, truncated = logged[:logMessageLimit], true + } + c.hub.log.LogAttrs(ctx, slog.LevelDebug, event, + slog.String("client", c.id), + slog.Int("bytes", len(data)), + slog.String("message", string(logged)), + slog.Bool("truncated", truncated), + ) +} + +// withTimeout returns ctx with the timeout applied, or ctx itself if the +// timeout is 0 or less (no limit). +func withTimeout( + ctx context.Context, timeout time.Duration, +) (context.Context, context.CancelFunc) { + if timeout <= 0 { + return ctx, func() {} + } + return context.WithTimeout(ctx, timeout) +} diff --git a/internal/websocket/client_test.go b/internal/websocket/client_test.go new file mode 100644 index 0000000..75a8641 --- /dev/null +++ b/internal/websocket/client_test.go @@ -0,0 +1,258 @@ +package websocket + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "io" + "log/slog" + "strings" + "sync" + "testing" + "time" + + "git.avsdev.uk/Go/job-search/internal/config" + ws "github.com/coder/websocket" +) + +// fakeConn is a msgConn that records how it was closed. Read blocks until its +// context is cancelled. +type fakeConn struct { + closed chan ws.StatusCode +} + +func newFakeConn() *fakeConn { return &fakeConn{closed: make(chan ws.StatusCode, 1)} } + +func (f *fakeConn) Read(ctx context.Context) (ws.MessageType, []byte, error) { + <-ctx.Done() + return 0, nil, ctx.Err() +} + +func (f *fakeConn) Write(context.Context, ws.MessageType, []byte) error { return nil } + +func (f *fakeConn) Ping(context.Context) error { return nil } + +func (f *fakeConn) Close(code ws.StatusCode, _ string) error { + f.closed <- code + return nil +} + +func (f *fakeConn) CloseNow() error { return nil } + +// newTestClient returns a client of a new hub on f, with no loops running, +// and its context. +func newTestClient(f *fakeConn) (*client, context.Context) { + h := NewHub(config.JSCServerWebSocket{}, discardLog) + ctx, cancel := context.WithCancel(context.Background()) + return newClient(h, f, "192.0.2.1", cancel), ctx +} + +// An envelope that can't be marshalled is dropped, not queued. +func TestSendEnvelopeBadPayload(t *testing.T) { + c, _ := newTestClient(newFakeConn()) + c.sendEnvelope(Envelope{Type: TypeEcho, Payload: json.RawMessage("{")}) + if len(c.send) != 0 { + t.Errorf("queued %d messages; want none", len(c.send)) + } +} + +// A client whose queue is full is closed as a slow consumer. +func TestSendEnvelopeSlowConsumer(t *testing.T) { + f := newFakeConn() + c, ctx := newTestClient(f) + for range clientSendBuffer { + c.sendEnvelope(Envelope{Type: TypePong}) + } + select { + case code := <-f.closed: + t.Fatalf("closed with %v before the queue was full", code) + default: + } + + c.sendEnvelope(Envelope{Type: TypePong}) + select { + case code := <-f.closed: + if code != ws.StatusPolicyViolation { + t.Errorf("closed with %v; want policy violation", code) + } + case <-time.After(time.Second): + t.Fatal("slow consumer not closed") + } + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("slow consumer not cancelled") + } +} + +// close closes the connection once, however often it is called. +func TestClientCloseOnce(t *testing.T) { + f := newFakeConn() + c, _ := newTestClient(f) + c.close(ws.StatusGoingAway, "bye") + c.close(ws.StatusNormalClosure, "again") + if code := <-f.closed; code != ws.StatusGoingAway { + t.Errorf("closed with %v; want going away", code) + } + select { + case code := <-f.closed: + t.Errorf("closed a second time with %v", code) + default: + } +} + +// The message rate limit applies only when it is enabled and has a rate. +func TestClientMessageRateLimit(t *testing.T) { + tests := []struct { + rule config.JSCRateLimitRule + limited bool + }{ + {config.JSCRateLimitRule{Enabled: false, PerSecond: 1, Burst: 1}, false}, + {config.JSCRateLimitRule{Enabled: true, PerSecond: 0, Burst: 1}, false}, + {config.JSCRateLimitRule{Enabled: true, PerSecond: 1, Burst: 1}, true}, + } + for _, tt := range tests { + h := NewHub(config.JSCServerWebSocket{MessageRateLimit: tt.rule}, discardLog) + c := newClient(h, newFakeConn(), "", func() {}) + c.limiter.Allow() + if limited := !c.limiter.Allow(); limited != tt.limited { + t.Errorf("rule %+v: second message limited = %v; want %v", tt.rule, limited, tt.limited) + } + } +} + +func TestMustRaw(t *testing.T) { + raw := json.RawMessage(`{"a":1}`) + tests := []struct { + name string + v any + want string + }{ + {"nil", nil, ""}, + {"raw message as is", raw, `{"a":1}`}, + {"marshalled", map[string]int{"b": 2}, `{"b":2}`}, + {"unmarshallable", func() {}, ""}, + } + for _, tt := range tests { + if got := string(mustRaw(tt.v)); got != tt.want { + t.Errorf("%s: mustRaw = %q; want %q", tt.name, got, tt.want) + } + } +} + +func TestWithTimeout(t *testing.T) { + ctx, cancel := withTimeout(context.Background(), 0) + defer cancel() + if _, ok := ctx.Deadline(); ok { + t.Error("timeout 0 set a deadline") + } + ctx, cancel = withTimeout(context.Background(), time.Minute) + defer cancel() + if _, ok := ctx.Deadline(); !ok { + t.Error("timeout 1m set no deadline") + } +} + +// syncBuffer is a bytes.Buffer safe for the client's goroutines to log to. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *syncBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +// entries returns the JSON log entries written so far. +func (b *syncBuffer) entries(t *testing.T) []map[string]any { + t.Helper() + b.mu.Lock() + defer b.mu.Unlock() + var entries []map[string]any + for line := range strings.Lines(b.buf.String()) { + var e map[string]any + if err := json.Unmarshal([]byte(line), &e); err != nil { + t.Fatalf("log line %q: %v", line, err) + } + entries = append(entries, e) + } + return entries +} + +// jsonLog returns a logger writing JSON entries at level and above to a +// buffer. +func jsonLog(level slog.Level) (*slog.Logger, *syncBuffer) { + buf := &syncBuffer{} + return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: level})), buf +} + +// At debug level, messages received from and sent to a client are logged. +func TestLogMessages(t *testing.T) { + log, buf := jsonLog(slog.LevelDebug) + h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, log) + peer, _ := pipeClient(t, h) + + ping := `{"type":"ping","id":"1"}` + go io.WriteString(peer, ping+"\n") + peer.SetReadDeadline(time.Now().Add(time.Second)) + pong, err := bufio.NewReader(peer).ReadString('\n') + if err != nil { + t.Fatalf("pong: %v", err) + } + pong = strings.TrimSuffix(pong, "\n") + + // The write is logged after it returns, so wait for its entry. + want := map[string]string{ + "message received": ping, + "message sent": pong, + } + deadline := time.Now().Add(time.Second) + for { + found := map[string]bool{} + for _, e := range buf.entries(t) { + if msg, ok := e["msg"].(string); ok && want[msg] == e["message"] { + found[msg] = true + if e["bytes"] != float64(len(want[msg])) || e["truncated"] != false { + t.Errorf("entry %v: bad bytes or truncated", e) + } + } + } + if len(found) == len(want) { + break + } + if time.Now().After(deadline) { + t.Fatalf("logged %v; want received ping and sent pong", buf.entries(t)) + } + time.Sleep(5 * time.Millisecond) + } +} + +// A long message is logged cut to logMessageLimit bytes; nothing is logged +// unless debug logging is on. +func TestLogMessageLimit(t *testing.T) { + log, buf := jsonLog(slog.LevelDebug) + h := NewHub(config.JSCServerWebSocket{}, log) + c := newClient(h, newFakeConn(), "", func() {}) + long := strings.Repeat("x", logMessageLimit+100) + c.logMessage(context.Background(), "message sent", []byte(long)) + + entries := buf.entries(t) + if len(entries) != 1 { + t.Fatalf("logged %d entries; want 1", len(entries)) + } + e := entries[0] + if e["message"] != long[:logMessageLimit] || e["bytes"] != float64(len(long)) || + e["truncated"] != true || e["client"] != c.id { + t.Errorf("entry = %v", e) + } + + log, buf = jsonLog(slog.LevelInfo) + c.hub = NewHub(config.JSCServerWebSocket{}, log) + c.logMessage(context.Background(), "message sent", []byte("x")) + if entries := buf.entries(t); len(entries) != 0 { + t.Errorf("logged %v at info level", entries) + } +} diff --git a/internal/websocket/http.go b/internal/websocket/http.go new file mode 100644 index 0000000..b28aab1 --- /dev/null +++ b/internal/websocket/http.go @@ -0,0 +1,44 @@ +package websocket + +import ( + "context" + "net/http" + + "git.avsdev.uk/Go/job-search/internal/middleware" + ws "github.com/coder/websocket" +) + +// ServeHTTP upgrades an http connection and runs the client until it closes. +func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) { + mode := ws.CompressionDisabled + if h.conf.EnableCompression { + mode = ws.CompressionContextTakeover + } + + // Accept enforces the Origin check: same host, or one of OriginPatterns. + // That is the defence against cross-site WebSocket hijacking. + conn, err := ws.Accept(w, r, &ws.AcceptOptions{ + CompressionMode: mode, + }) + if err != nil { + h.log.Debug("upgrade failed", "error", err) + return // Accept has already written the HTTP error. + } + defer conn.CloseNow() //nolint:errcheck // finished with; it may already be closed + conn.SetReadLimit(int64(h.maxMessageSize())) + + ip := middleware.ClientIP(r) + ctx, cancel := context.WithCancel(r.Context()) + c := newClient(h, conn, ip, cancel) + + if !h.register(c) { + cancel() + conn.Close(ws.StatusGoingAway, "server shutting down") //nolint:errcheck,gosec // turned away + return + } + defer h.unregister(c) + + h.log.Info("client connected", "client", c.id, "ip", ip) + c.run(ctx) + h.log.Info("client disconnected", "client", c.id, "ip", ip) +} diff --git a/internal/websocket/http_test.go b/internal/websocket/http_test.go new file mode 100644 index 0000000..9b2f88e --- /dev/null +++ b/internal/websocket/http_test.go @@ -0,0 +1,183 @@ +package websocket + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "git.avsdev.uk/Go/job-search/internal/config" + ws "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" +) + +// wsHub serves h's WebSocket clients on a test HTTP server and returns its +// ws:// URL. +func wsHub(t *testing.T, h *Hub) string { + t.Helper() + srv := httptest.NewServer(h) + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + h.Shutdown(ctx) + srv.Close() + }) + return "ws" + strings.TrimPrefix(srv.URL, "http") +} + +// dialWS connects a WebSocket client to url and reads its welcome. +func dialWS(t *testing.T, ctx context.Context, url string) *ws.Conn { + t.Helper() + conn, _, err := ws.Dial(ctx, url, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + t.Cleanup(func() { conn.CloseNow() }) + var e Envelope + if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypeWelcome { + t.Fatalf("want welcome, got %+v, %v", e, err) + } + return conn +} + +// A client upgraded over HTTP gets replies. A max_message_size of 0 uses the +// default limit rather than refusing every message. +func TestHTTPMessages(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, discardLog) + url := wsHub(t, h) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + conn := dialWS(t, ctx, url) + + if err := wsjson.Write(ctx, conn, Envelope{Type: TypePing, ID: "1"}); err != nil { + t.Fatal(err) + } + var e Envelope + if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypePong || e.ID != "1" { + t.Fatalf("want pong 1, got %+v, %v", e, err) + } +} + +// One Push reaches WebSocket and plain TCP clients of the same hub. +func TestPushReachesBothTransports(t *testing.T) { + h, addr := tcpHub(t, nil) + url := wsHub(t, h) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + conn := dialWS(t, ctx, url) + c := dialTCP(t, addr) + c.read() // welcome + waitCount(t, h, 2) + + h.Push("item.created", map[string]string{"id": "x"}) + var e Envelope + if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != "item.created" { + t.Fatalf("websocket: want item.created, got %+v, %v", e, err) + } + if e := c.read(); e.Type != "item.created" { + t.Fatalf("tcp: want item.created, got %+v", e) + } +} + +// A payload that can't be marshalled is not sent. +func TestPushBadPayload(t *testing.T) { + h, addr := tcpHub(t, nil) + c := dialTCP(t, addr) + c.read() // welcome + waitCount(t, h, 1) + + h.Push("bad", func() {}) + h.Push("good", nil) + if e := c.read(); e.Type != "good" || e.Payload != nil { + t.Fatalf("want good with no payload, got %+v", e) + } +} + +// A binary frame closes the connection with 1003 (unsupported data). +func TestHTTPBinaryFrame(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, discardLog) + url := wsHub(t, h) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + conn := dialWS(t, ctx, url) + + if err := conn.Write(ctx, ws.MessageBinary, []byte{1, 2, 3}); err != nil { + t.Fatal(err) + } + _, _, err := conn.Read(ctx) + if ws.CloseStatus(err) != ws.StatusUnsupportedData { + t.Fatalf("read after binary frame = %v; want close 1003", err) + } +} + +// Messages over the message rate limit get a "rate limited" error and are not +// handled; the connection stays open. +func TestHTTPMessageRateLimit(t *testing.T) { + conf := config.JSCServerWebSocket{ + WriteTimeout: time.Second, + MessageRateLimit: config.JSCRateLimitRule{Enabled: true, PerSecond: 1, Burst: 1}, + } + h := NewHub(conf, discardLog) + url := wsHub(t, h) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + conn := dialWS(t, ctx, url) + + for _, id := range []string{"1", "2"} { + if err := wsjson.Write(ctx, conn, Envelope{Type: TypePing, ID: id}); err != nil { + t.Fatal(err) + } + } + var e Envelope + if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypePong || e.ID != "1" { + t.Fatalf("first: want pong 1, got %+v, %v", e, err) + } + err := wsjson.Read(ctx, conn, &e) + if err != nil || e.Type != TypeError || e.Error != "rate limited" { + t.Fatalf("second: want rate limited, got %+v, %v", e, err) + } + if h.Count() != 1 { + t.Errorf("Count = %d; want the client kept", h.Count()) + } +} + +// A request that isn't a WebSocket upgrade is refused by the upgrade and no +// client is registered. +func TestHTTPNotAnUpgrade(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{}, discardLog) + srv := httptest.NewServer(h) + defer srv.Close() + res, err := http.Get(srv.URL) + if err != nil { + t.Fatal(err) + } + res.Body.Close() + if res.StatusCode < 400 { + t.Errorf("status = %d; want an error", res.StatusCode) + } + if h.Count() != 0 { + t.Errorf("Count = %d; want 0", h.Count()) + } +} + +// After Shutdown, an upgrade succeeds but the client is closed at once with +// 1001 (going away). +func TestHTTPRefusedAfterShutdown(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{}, discardLog) + url := wsHub(t, h) + if err := h.Shutdown(context.Background()); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + conn, _, err := ws.Dial(ctx, url, nil) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer conn.CloseNow() + if _, _, err := conn.Read(ctx); ws.CloseStatus(err) != ws.StatusGoingAway { + t.Fatalf("read = %v; want close 1001", err) + } +} diff --git a/internal/websocket/hub.go b/internal/websocket/hub.go new file mode 100644 index 0000000..ca470df --- /dev/null +++ b/internal/websocket/hub.go @@ -0,0 +1,158 @@ +package websocket + +import ( + "context" + "encoding/json" + "log/slog" + "sync" + + "git.avsdev.uk/Go/job-search/internal/config" + ws "github.com/coder/websocket" +) + +// Hub tracks connected clients, WebSocket (ServeHTTP) and plain TCP +// (ServeTCP) alike, and sends pushed messages to all of them. +type Hub struct { + conf config.JSCServerWebSocket + log *slog.Logger + + m sync.RWMutex + clients map[*client]struct{} + closed bool + wg sync.WaitGroup +} + +// defaultMaxMessageSize is the largest message read from a client when +// max_message_size is not set. +const defaultMaxMessageSize = 32 << 10 + +// NewHub creates a hub. +func NewHub(conf config.JSCServerWebSocket, log *slog.Logger) *Hub { + return &Hub{ + conf: conf, + log: log.With("component", "ws"), + clients: make(map[*client]struct{}), + } +} + +// Push sends a message of type msgType to every client. The payload is +// marshalled once and the same bytes are queued for every client; a client +// whose queue is full is closed as a slow consumer. A payload that can't be +// marshalled is logged and nothing is sent. +func (h *Hub) Push(msgType string, payload any) { + var rawPayload json.RawMessage + if payload != nil { + var err error + if rawPayload, err = json.Marshal(payload); err != nil { + h.log.Error("marshal push payload", "type", msgType, "error", err) + return + } + } + + msg, err := json.Marshal(Envelope{ + Type: msgType, + Payload: rawPayload, + }) + if err != nil { + h.log.Error("marshal push", "type", msgType, "error", err) + return + } + + var slow []*client + h.m.RLock() + for client := range h.clients { + if !client.trySend(msg) { + slow = append(slow, client) + } + } + h.m.RUnlock() + + for _, client := range slow { + h.log.Warn("dropping slow client", "client", client.id, "ip", client.ip) + go client.close(ws.StatusPolicyViolation, "slow consumer") + } +} + +// Count returns the number of connected clients. +func (h *Hub) Count() int { + h.m.RLock() + defer h.m.RUnlock() + return len(h.clients) +} + +// Shutdown closes every connection with 1001 (going away) and waits for the +// handlers to return. If ctx expires first, the remaining connections are +// closed at once and ctx's error is returned. New connections are refused +// afterwards. +func (h *Hub) Shutdown(ctx context.Context) error { + h.m.Lock() + h.closed = true + clientShutdowns := make([]*client, 0, len(h.clients)) + for client := range h.clients { + clientShutdowns = append(clientShutdowns, client) + } + h.m.Unlock() + + if len(clientShutdowns) == 0 { + // Nothing to close. Clients that just left are only finishing their + // unregister, so this wait is short and ctx doesn't apply. + h.wg.Wait() + return nil + } + + for _, client := range clientShutdowns { + go client.close(ws.StatusGoingAway, "server shutting down") + } + + done := make(chan struct{}) + go func() { + h.wg.Wait() + close(done) + }() + + select { + case <-done: + return nil + case <-ctx.Done(): + // Both may be ready; a finished shutdown is not a timeout. + select { + case <-done: + return nil + default: + } + for _, client := range clientShutdowns { + client.conn.CloseNow() //nolint:errcheck,gosec // out of time; dropping it + } + return ctx.Err() + } +} + +// maxMessageSize is the largest message read from a client: max_message_size, +// or defaultMaxMessageSize if that is 0 or less. +func (h *Hub) maxMessageSize() int { + if h.conf.MaxMessageSize <= 0 { + return defaultMaxMessageSize + } + return int(h.conf.MaxMessageSize) +} + +// register adds a client to the hub. Hub clients receive all push messages and +// have their lifecycles tracked by the hub. +func (h *Hub) register(c *client) bool { + h.m.Lock() + defer h.m.Unlock() + if h.closed { + return false + } + h.clients[c] = struct{}{} + h.wg.Add(1) + return true +} + +// unregister removes a client from the hub. +func (h *Hub) unregister(c *client) { + h.m.Lock() + delete(h.clients, c) + h.m.Unlock() + h.wg.Done() +} diff --git a/internal/websocket/hub_test.go b/internal/websocket/hub_test.go new file mode 100644 index 0000000..1ee2108 --- /dev/null +++ b/internal/websocket/hub_test.go @@ -0,0 +1,93 @@ +package websocket + +import ( + "bufio" + "context" + "errors" + "net" + "testing" + "time" + + "git.avsdev.uk/Go/job-search/internal/config" +) + +// pipeClient runs a plain TCP client of h on an in-memory pipe and reads its +// welcome. It returns the peer's end, which the test stops reading from to +// make the client's writes block, and a channel closed when ServeTCP returns. +func pipeClient(t *testing.T, h *Hub) (net.Conn, chan struct{}) { + t.Helper() + server, peer := net.Pipe() + t.Cleanup(func() { server.Close(); peer.Close() }) + done := make(chan struct{}) + go func() { h.ServeTCP(context.Background(), server); close(done) }() + peer.SetReadDeadline(time.Now().Add(time.Second)) + if _, err := bufio.NewReader(peer).ReadString('\n'); err != nil { + t.Fatalf("welcome: %v", err) + } + waitCount(t, h, 1) + return peer, done +} + +// If ctx expires before the clients have closed, Shutdown closes the rest at +// once and returns ctx's error. +func TestShutdownTimeout(t *testing.T) { + // No write timeout: the client's goodbye blocks until it is cut short. + h := NewHub(config.JSCServerWebSocket{}, discardLog) + _, done := pipeClient(t, h) + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if err := h.Shutdown(ctx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Shutdown = %v; want context.DeadlineExceeded", err) + } + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("client not closed after the shutdown timeout") + } +} + +// Shutdown with no clients returns at once, even with an expired ctx. +func TestShutdownNoClients(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{}, discardLog) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + for range 20 { + if err := h.Shutdown(ctx); err != nil { + t.Fatalf("Shutdown = %v; want nil", err) + } + } +} + +// A ping that can't be sent within the pong timeout drops the client. +func TestPingFailureDropsClient(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{ + PingInterval: 10 * time.Millisecond, + PongTimeout: 20 * time.Millisecond, + }, discardLog) + _, done := pipeClient(t, h) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("client not dropped after a failed ping") + } + waitCount(t, h, 0) +} + +func TestMaxMessageSize(t *testing.T) { + tests := []struct { + conf config.JSCServerWebSocket + want int + }{ + {config.JSCServerWebSocket{}, defaultMaxMessageSize}, + {config.JSCServerWebSocket{MaxMessageSize: -1}, defaultMaxMessageSize}, + {config.JSCServerWebSocket{MaxMessageSize: 100}, 100}, + } + for _, tt := range tests { + h := NewHub(tt.conf, discardLog) + if got := h.maxMessageSize(); got != tt.want { + t.Errorf("maxMessageSize(%d) = %d; want %d", tt.conf.MaxMessageSize, got, tt.want) + } + } +} diff --git a/internal/websocket/messages.go b/internal/websocket/messages.go new file mode 100644 index 0000000..70da3ca --- /dev/null +++ b/internal/websocket/messages.go @@ -0,0 +1,64 @@ +// Package websocket implements the hub that pushes messages to connected +// clients: WebSocket clients upgraded by Hub.ServeHTTP (the /ws endpoint) and +// plain TCP clients run by Hub.ServeTCP, which send the same messages as +// newline-delimited JSON (see tcp.go). Package tcpserver runs a TCP server +// for those clients. +// +// Every message is a JSON Envelope: a WebSocket text frame, or one line on +// TCP. +// +// Client -> server: +// +// {"type":"ping","id":"1"} +// {"type":"echo","id":"2","payload":{...}} +// +// Server -> client: +// +// {"type":"welcome","payload":{"clientId":"..."}} +// {"type":"pong","id":"1","payload":{"time":"..."}} +// {"type":"echo","id":"2","payload":{...}} +// {"type":"error","id":"3","error":"..."} +// {"type":"item.created","payload":{...}} +// +// Any other type from a client gets an "unknown message type" error. The last +// message is one sent by Hub.Push: its type is the pushed message type. +package websocket + +import "encoding/json" + +// Envelope is the wire format for every message in both directions. +type Envelope struct { + Type string `json:"type"` + ID string `json:"id,omitempty"` + Payload json.RawMessage `json:"payload,omitempty"` + Error string `json:"error,omitempty"` +} + +// Envelope types. TypeEvent, TypeRequest and TypeAck are reserved: nothing +// sends or handles them yet. +const ( + TypeWelcome = "welcome" + TypePing = "ping" + TypePong = "pong" + TypeEcho = "echo" + TypeEvent = "event" + TypeRequest = "request" + TypeAck = "ack" + TypeError = "error" +) + +// mustRaw returns v as raw JSON: v itself if it is already json.RawMessage, +// nil if v is nil or can't be marshalled. +func mustRaw(v any) json.RawMessage { + if v == nil { + return nil + } + if raw, ok := v.(json.RawMessage); ok { + return raw + } + b, err := json.Marshal(v) + if err != nil { + return nil + } + return b +} diff --git a/internal/websocket/tcp.go b/internal/websocket/tcp.go new file mode 100644 index 0000000..f06b7cf --- /dev/null +++ b/internal/websocket/tcp.go @@ -0,0 +1,261 @@ +package websocket + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "net" + "net/netip" + "strings" + "sync" + "time" + + ws "github.com/coder/websocket" +) + +// Plain TCP clients join the hub as if they were WebSocket clients. Instead +// of WebSocket frames, each message is one line of JSON (an Envelope) +// terminated by "\n" in both directions, so a client can be as simple as +// netcat: +// +// $ nc localhost 9393 +// {"type":"welcome","payload":{"clientId":"..."}} +// {"type":"ping","id":"1"} +// {"type":"pong","id":"1","payload":{"time":"..."}} +// +// A trailing "\r" is ignored and blank lines are skipped. A line longer than +// the websocket max_message_size closes the connection, as an oversized frame +// does. When the server closes a connection (shutdown, slow consumer, a line +// too long) it sends an error envelope with the reason first. +// +// Plain TCP has no ping frames, so liveness uses TCP keepalives probing at the +// ping interval; Ping itself writes a blank line, which clients ignore, so a +// dead peer is found when a write fails. +// +// There is no HTTP request, so there is no CSRF or Origin check and no proxy +// handling: the client's address is the TCP peer. Only listen where clients +// are trusted (e.g. localhost), or wrap the listener (e.g. with PROXY protocol +// or TLS) before passing its connections to ServeTCP. ServeTCP doesn't limit +// connections; package tcpserver applies websocket.connection_limit. + +// ServeTCP runs a plain TCP connection as a hub client until it closes, as +// ServeHTTP does for a WebSocket upgrade. Cancelling ctx drops the client +// without a message; use Shutdown to close clients gracefully. The +// connection is closed when ServeTCP returns. +func (h *Hub) ServeTCP(ctx context.Context, newConn net.Conn) { + if tcpSocket, ok := newConn.(*net.TCPConn); ok && h.conf.PingInterval > 0 { + tcpSocket.SetKeepAliveConfig(net.KeepAliveConfig{ //nolint:errcheck,gosec // best effort + Enable: true, + Idle: h.conf.PingInterval, + Interval: h.conf.PingInterval, + Count: 3, + }) + } + + ip := "" + if addrPort, err := netip.ParseAddrPort(newConn.RemoteAddr().String()); err == nil { + ip = addrPort.Addr().Unmap().String() + } + + conn := newTCPConn(newConn, h.maxMessageSize(), h.conf.WriteTimeout) + + ctx, cancel := context.WithCancel(ctx) + client := newClient(h, conn, ip, cancel) + + if !h.register(client) { + cancel() + conn.Close(ws.StatusGoingAway, "server shutting down") //nolint:errcheck,gosec // turned away + return + } + defer h.unregister(client) + + h.log.Info("client connected", "client", client.id, "ip", ip, "transport", "tcp") + client.run(ctx) + conn.CloseNow() //nolint:errcheck,gosec // finished with; it may already be closed + h.log.Info("client disconnected", "client", client.id, "ip", ip, "transport", "tcp") +} + +// ---------- tcpConn ---------- + +// tcpConn adapts a net.Conn carrying newline-delimited JSON to the msgConn a +// client uses. Read is only called from the client's read loop; writes may +// come from the write loop, Ping and Close concurrently, so they are +// serialised. +type tcpConn struct { + conn net.Conn + scan *bufio.Scanner + closeTimeout time.Duration // caps Close's message; 0 or less for none + + wm sync.Mutex // serialises writes + + closeOnce sync.Once // records the first close + closed chan struct{} // closed by the first Close or CloseNow + closeErr ws.CloseError // why the connection was closed +} + +// newTCPConn wraps conn. maxLine caps the length of a line read +// (defaultMaxMessageSize if 0 or less) and closeTimeout how long Close spends +// sending its message. +func newTCPConn( + conn net.Conn, maxLine int, closeTimeout time.Duration, +) *tcpConn { + if maxLine <= 0 { + maxLine = defaultMaxMessageSize + } + + scan := bufio.NewScanner(conn) + // The buffer must hold the line and its "\n". + scan.Buffer(make([]byte, 0, min(maxLine+1, 4096)), maxLine+1) + + return &tcpConn{ + conn: conn, + scan: scan, + closeTimeout: closeTimeout, + closed: make(chan struct{}), + } +} + +// Read returns the next non-blank line as a text message. Cancelling ctx +// interrupts it. The end of the stream is reported as a normal closure, and a +// read after Close as the close status Close sent, as a *websocket.Conn +// reports them, so the client treats them as a finished close handshake. +func (t *tcpConn) Read(ctx context.Context) (ws.MessageType, []byte, error) { + stop := context.AfterFunc(ctx, func() { + t.conn.SetReadDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the read + }) + defer stop() + + for t.scan.Scan() { + line := strings.TrimSuffix(t.scan.Text(), "\r") + if strings.TrimSpace(line) == "" { + continue + } + return ws.MessageText, []byte(line), nil + } + + err := t.scan.Err() + switch { + case ctx.Err() != nil: + return 0, nil, ctx.Err() + case t.isClosed(): + return 0, nil, t.closeErr + case err == nil || errors.Is(err, io.EOF): + return 0, nil, ws.CloseError{ + Code: ws.StatusNormalClosure, + Reason: "connection closed", + } + case errors.Is(err, bufio.ErrTooLong): + // As a *websocket.Conn closes with 1009 for an oversized message. + // The client may not see the reason: closing with its unread data + // still buffered makes the connection reset. + t.Close(ws.StatusMessageTooBig, "message too long") //nolint:errcheck,gosec // dropped anyway + return 0, nil, t.closeErr + } + return 0, nil, err +} + +// Write sends payload as one line. Only text messages are supported. +func (t *tcpConn) Write( + ctx context.Context, messageType ws.MessageType, payload []byte, +) error { + if messageType != ws.MessageText { + return errUnsupportedData + } + line := make([]byte, 0, len(payload)+1) + line = append(append(line, payload...), '\n') + return t.write(ctx, line) +} + +// Ping writes a blank line, which clients ignore, so a dead peer shows up as a +// failed write. TCP keepalives (see ServeTCP) do the rest. +func (t *tcpConn) Ping(ctx context.Context) error { + return t.write(ctx, []byte("\n")) +} + +// Close sends an error envelope giving the reason, then closes the +// connection. Only the first Close or CloseNow does anything; later calls +// return nil. A CloseNow while Close is sending its message cuts it short. +func (t *tcpConn) Close(code ws.StatusCode, reason string) error { + if !t.markClosed(ws.CloseError{Code: code, Reason: reason}) { + return nil + } + + ctx := context.Background() + if t.closeTimeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, t.closeTimeout) + defer cancel() + } + + message, err := json.Marshal(Envelope{ + Type: TypeError, + Error: reason, + }) + if err == nil { + t.write(ctx, append(message, '\n')) //nolint:errcheck,gosec // best effort, closing anyway + } + + if err := t.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) { + return err + } + return nil +} + +// CloseNow closes the connection without a message. It doesn't wait for a +// Close that is sending its message. +func (t *tcpConn) CloseNow() error { + t.markClosed(ws.CloseError{Code: ws.StatusAbnormalClosure}) + if err := t.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) { + return err + } + return nil +} + +// markClosed records why the connection is being closed and reports whether +// this is the first close. It never blocks, so a Close sending its message +// can't hold up CloseNow. +func (t *tcpConn) markClosed(why ws.CloseError) bool { + first := false + t.closeOnce.Do(func() { + t.closeErr = why + close(t.closed) + first = true + }) + return first +} + +// isClosed reports whether Close or CloseNow has been called. +func (t *tcpConn) isClosed() bool { + select { + case <-t.closed: + return true + default: + return false + } +} + +// write writes data whole, giving up at ctx's deadline or when ctx is +// cancelled. +func (t *tcpConn) write(ctx context.Context, data []byte) error { + t.wm.Lock() + defer t.wm.Unlock() + + deadline, _ := ctx.Deadline() // zero means none + if err := t.conn.SetWriteDeadline(deadline); err != nil { + return err + } + + stop := context.AfterFunc(ctx, func() { + t.conn.SetWriteDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the write + }) + defer stop() + + _, err := t.conn.Write(data) + if err != nil && ctx.Err() != nil { + return ctx.Err() + } + + return err +} diff --git a/internal/websocket/tcp_test.go b/internal/websocket/tcp_test.go new file mode 100644 index 0000000..83bee05 --- /dev/null +++ b/internal/websocket/tcp_test.go @@ -0,0 +1,405 @@ +package websocket + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net" + "strings" + "syscall" + "testing" + "time" + + "git.avsdev.uk/Go/job-search/internal/config" + ws "github.com/coder/websocket" +) + +var discardLog = slog.New(slog.NewTextHandler(io.Discard, nil)) + +// tcpHub starts a hub serving plain TCP clients, via ServeTCP, on a local +// port and returns it with the address. It is shut down when the test ends. +func tcpHub(t *testing.T, mutate func(*config.JSCServerWebSocket)) (*Hub, string) { + t.Helper() + conf := config.JSCServerWebSocket{ + MaxMessageSize: 1 << 10, + WriteTimeout: time.Second, + } + if mutate != nil { + mutate(&conf) + } + h := NewHub(conf, discardLog) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return + } + go h.ServeTCP(ctx, conn) + } + }() + t.Cleanup(func() { + ln.Close() + sctx, scancel := context.WithTimeout(context.Background(), time.Second) + defer scancel() + h.Shutdown(sctx) + cancel() + }) + return h, ln.Addr().String() +} + +// tcpClient is a plain TCP test client reading newline-delimited envelopes. +type tcpClient struct { + t *testing.T + conn net.Conn + r *bufio.Reader +} + +func dialTCP(t *testing.T, addr string) *tcpClient { + t.Helper() + conn, err := net.Dial("tcp", addr) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { conn.Close() }) + return &tcpClient{t: t, conn: conn, r: bufio.NewReader(conn)} +} + +func (c *tcpClient) send(line string) { + c.t.Helper() + if _, err := io.WriteString(c.conn, line); err != nil { + c.t.Fatal(err) + } +} + +// read returns the next envelope, failing the test after a second. +func (c *tcpClient) read() Envelope { + c.t.Helper() + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + line, err := c.r.ReadString('\n') + if err != nil { + c.t.Fatalf("read: %v (got %q)", err, line) + } + var env Envelope + if err := json.Unmarshal([]byte(line), &env); err != nil { + c.t.Fatalf("bad envelope %q: %v", line, err) + } + return env +} + +// waitClosed checks the server closes the connection, skipping any lines +// sent first. A reset counts as closed: the server closing with unread data +// from the client resets the connection. +func (c *tcpClient) waitClosed() { + c.t.Helper() + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + if _, err := io.ReadAll(c.r); err != nil && !errors.Is(err, syscall.ECONNRESET) { + c.t.Fatalf("connection not closed: %v", err) + } +} + +func waitCount(t *testing.T, h *Hub, want int) { + t.Helper() + for deadline := time.Now().Add(time.Second); h.Count() != want; { + if time.Now().After(deadline) { + t.Fatalf("Count = %d; want %d", h.Count(), want) + } + time.Sleep(5 * time.Millisecond) + } +} + +// ---------- ServeTCP ---------- + +func TestTCPMessages(t *testing.T) { + h, addr := tcpHub(t, nil) + c := dialTCP(t, addr) + + if e := c.read(); e.Type != TypeWelcome { + t.Fatalf("want welcome, got %+v", e) + } + waitCount(t, h, 1) + + // Blank lines are skipped and a trailing \r is ignored. + c.send("\n \r\n{\"type\":\"ping\",\"id\":\"1\"}\r\n") + if e := c.read(); e.Type != TypePong || e.ID != "1" { + t.Fatalf("want pong 1, got %+v", e) + } + c.send(`{"type":"echo","id":"2","payload":{"a":1}}` + "\n") + if e := c.read(); e.Type != TypeEcho || e.ID != "2" || string(e.Payload) != `{"a":1}` { + t.Fatalf("want echo 2, got %+v", e) + } + c.send("not json\n") + if e := c.read(); e.Type != TypeError || e.Error != "invalid JSON envelope" { + t.Fatalf("want invalid JSON error, got %+v", e) + } + c.send(`{"type":"bogus","id":"3"}` + "\n") + if e := c.read(); e.Type != TypeError || e.ID != "3" { + t.Fatalf("want unknown type error, got %+v", e) + } + + // Pushed messages reach TCP clients. + h.Push("item.created", map[string]string{"id": "x"}) + if e := c.read(); e.Type != "item.created" || string(e.Payload) != `{"id":"x"}` { + t.Fatalf("want item.created, got %+v", e) + } +} + +// With no write timeout, messages are still sent: a zero timeout must not +// mean one that has already expired. +func TestTCPNoWriteTimeout(t *testing.T) { + _, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { c.WriteTimeout = 0 }) + c := dialTCP(t, addr) + if e := c.read(); e.Type != TypeWelcome { + t.Fatalf("want welcome, got %+v", e) + } + c.send(`{"type":"ping","id":"1"}` + "\n") + if e := c.read(); e.Type != TypePong { + t.Fatalf("want pong, got %+v", e) + } +} + +func TestTCPClientDisconnect(t *testing.T) { + h, addr := tcpHub(t, nil) + c := dialTCP(t, addr) + c.read() // welcome + waitCount(t, h, 1) + c.conn.Close() + waitCount(t, h, 0) +} + +// A line longer than max_message_size drops the connection. +func TestTCPLineTooLong(t *testing.T) { + h, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { c.MaxMessageSize = 64 }) + c := dialTCP(t, addr) + c.read() // welcome + c.send(`{"type":"echo","payload":"` + strings.Repeat("x", 100) + `"}` + "\n") + c.waitClosed() + waitCount(t, h, 0) +} + +// Shutdown tells TCP clients why they are being closed, then closes them, +// and refuses new ones. +func TestTCPShutdown(t *testing.T) { + h, addr := tcpHub(t, nil) + c := dialTCP(t, addr) + c.read() // welcome + waitCount(t, h, 1) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := h.Shutdown(ctx); err != nil { + t.Fatalf("Shutdown: %v", err) + } + if e := c.read(); e.Type != TypeError || e.Error != "server shutting down" { + t.Fatalf("want shutdown error, got %+v", e) + } + c.waitClosed() + + late := dialTCP(t, addr) + if e := late.read(); e.Type != TypeError || e.Error != "server shutting down" { + t.Fatalf("late client: want shutdown error, got %+v", e) + } + late.waitClosed() +} + +// Ping writes a blank line to an idle client. The client's messages, which +// also count as activity, can arrive while pings run. +func TestTCPPing(t *testing.T) { + _, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { + c.PingInterval = 20 * time.Millisecond + c.PongTimeout = time.Second + }) + c := dialTCP(t, addr) + c.read() // welcome + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + line, err := c.r.ReadString('\n') + if err != nil || line != "\n" { + t.Fatalf("want a blank ping line, got %q, %v", line, err) + } + for i := range 5 { + c.send(`{"type":"ping"}` + "\n") + for { + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + line, err := c.r.ReadString('\n') + if err != nil { + t.Fatalf("message %d: %v", i, err) + } + if line != "\n" { + break + } + } + time.Sleep(10 * time.Millisecond) + } +} + +// A client that doesn't read is dropped as a slow consumer once its queue is +// full. net.Pipe has no buffer, so writes to it block until the peer reads. +func TestTCPSlowConsumer(t *testing.T) { + h := NewHub(config.JSCServerWebSocket{WriteTimeout: 50 * time.Millisecond}, discardLog) + server, peer := net.Pipe() + t.Cleanup(func() { server.Close(); peer.Close() }) + done := make(chan struct{}) + go func() { h.ServeTCP(context.Background(), server); close(done) }() + waitCount(t, h, 1) + + for range clientSendBuffer + 1 { + h.Push("item.created", nil) + } + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("slow client not dropped") + } + waitCount(t, h, 0) +} + +// A pong timeout of 0 means no timeout, not one that has already expired: +// idle clients are pinged and kept. +func TestTCPNoPongTimeout(t *testing.T) { + h, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { + c.PingInterval = 10 * time.Millisecond + c.PongTimeout = 0 + }) + c := dialTCP(t, addr) + c.read() // welcome + for range 3 { + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + if line, err := c.r.ReadString('\n'); err != nil || line != "\n" { + t.Fatalf("want a blank ping line, got %q, %v", line, err) + } + } + waitCount(t, h, 1) +} + +// A max_message_size of 0 uses the default limit. +func TestTCPDefaultMessageSize(t *testing.T) { + _, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { c.MaxMessageSize = 0 }) + c := dialTCP(t, addr) + c.read() // welcome + c.send(`{"type":"echo","id":"1","payload":"` + strings.Repeat("x", 4096) + `"}` + "\n") + if e := c.read(); e.Type != TypeEcho || e.ID != "1" { + t.Fatalf("want echo, got %+v", e) + } +} + +// ---------- tcpConn ---------- + +// pipeConn returns a tcpConn on one end of an in-memory pipe, and the other +// end. Writes on a pipe block until the other end reads. +func pipeConn(t *testing.T, closeTimeout time.Duration) (*tcpConn, net.Conn) { + t.Helper() + server, peer := net.Pipe() + t.Cleanup(func() { server.Close(); peer.Close() }) + return newTCPConn(server, 0, closeTimeout), peer +} + +func TestTCPConnRead(t *testing.T) { + tc, peer := pipeConn(t, time.Second) + go func() { + io.WriteString(peer, "\r\n{\"a\":1}\r\n") + peer.Close() + }() + typ, data, err := tc.Read(context.Background()) + if err != nil || typ != ws.MessageText || string(data) != `{"a":1}` { + t.Fatalf("Read = %v, %q, %v", typ, data, err) + } + // The peer closing is a normal closure. + _, _, err = tc.Read(context.Background()) + if ws.CloseStatus(err) != ws.StatusNormalClosure { + t.Fatalf("Read at EOF = %v; want a normal closure", err) + } +} + +func TestTCPConnReadCancelled(t *testing.T) { + tc, _ := pipeConn(t, time.Second) + ctx, cancel := context.WithCancel(context.Background()) + time.AfterFunc(10*time.Millisecond, cancel) + if _, _, err := tc.Read(ctx); !errors.Is(err, context.Canceled) { + t.Fatalf("Read = %v; want context.Canceled", err) + } +} + +func TestTCPConnWriteBinary(t *testing.T) { + tc, _ := pipeConn(t, time.Second) + err := tc.Write(context.Background(), ws.MessageBinary, []byte("x")) + if !errors.Is(err, errUnsupportedData) { + t.Fatalf("Write(binary) = %v; want errUnsupportedData", err) + } +} + +// Close sends its reason and later reads report its status; a second Close +// or a CloseNow does nothing more. +func TestTCPConnClose(t *testing.T) { + tc, peer := pipeConn(t, time.Second) + got := make(chan string, 1) + go func() { + line, _ := bufio.NewReader(peer).ReadString('\n') + got <- line + }() + if err := tc.Close(ws.StatusGoingAway, "bye"); err != nil { + t.Fatalf("Close: %v", err) + } + if line := <-got; line != `{"type":"error","error":"bye"}`+"\n" { + t.Errorf("peer got %q", line) + } + _, _, err := tc.Read(context.Background()) + var ce ws.CloseError + if !errors.As(err, &ce) || ce.Code != ws.StatusGoingAway || ce.Reason != "bye" { + t.Errorf("Read after Close = %v; want going away, bye", err) + } + if err := tc.Close(ws.StatusNormalClosure, "again"); err != nil { + t.Errorf("second Close: %v", err) + } + if err := tc.CloseNow(); err != nil { + t.Errorf("CloseNow after Close: %v", err) + } +} + +// CloseNow doesn't wait for a Close stuck sending its message to a peer that +// isn't reading, and cuts it short. +func TestTCPConnCloseNowDuringClose(t *testing.T) { + tc, _ := pipeConn(t, 0) // no timeout: Close would wait for ever + closed := make(chan error, 1) + go func() { closed <- tc.Close(ws.StatusGoingAway, "bye") }() + for !tc.isClosed() { + time.Sleep(time.Millisecond) + } + + done := make(chan error, 1) + go func() { done <- tc.CloseNow() }() + select { + case err := <-done: + if err != nil { + t.Errorf("CloseNow: %v", err) + } + case <-time.After(time.Second): + t.Fatal("CloseNow waited for Close") + } + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("Close not cut short by CloseNow") + } +} + +// Close gives up sending its message after the close timeout. +func TestTCPConnCloseTimeout(t *testing.T) { + tc, _ := pipeConn(t, 20*time.Millisecond) + done := make(chan error, 1) + go func() { done <- tc.Close(ws.StatusGoingAway, "bye") }() + select { + case err := <-done: + if err != nil { + t.Errorf("Close: %v", err) + } + case <-time.After(time.Second): + t.Fatal("Close ignored its timeout") + } +} diff --git a/internal/websocket/tcpserver/server.go b/internal/websocket/tcpserver/server.go new file mode 100644 index 0000000..ecce52f --- /dev/null +++ b/internal/websocket/tcpserver/server.go @@ -0,0 +1,182 @@ +// Package tcpserver runs a TCP server whose clients join a websocket.Hub +// directly, without an HTTP upgrade: each connection carries newline-delimited +// JSON envelopes (see websocket.Hub.ServeTCP). It can serve a hub on its own, +// or provide a hub's clients in tests. +package tcpserver + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "net" + "net/netip" + "strconv" + "time" + + "git.avsdev.uk/Go/job-search/internal/config" + "git.avsdev.uk/Go/job-search/internal/middleware" + "git.avsdev.uk/Go/job-search/internal/websocket" +) + +// rejectTimeout caps how long telling a rate limited client why it is being +// turned away may take. +const rejectTimeout = time.Second + +// Server listens on the configured host and port and runs each connection as +// a client of its hub. New connections from each client IP are limited by +// websocket.connection_limit. +type Server struct { + conf *config.JSCServer + log *slog.Logger + hub *websocket.Hub +} + +// New builds a server for hub's plain TCP clients. Run starts it; it shuts +// hub down when it stops, so hub should not be shared with another server. +// Push to hub to reach the server's clients. +func New(conf *config.JSCServer, hub *websocket.Hub, log *slog.Logger) *Server { + return &Server{ + conf: conf, + log: log, + hub: hub, + } +} + +// Run listens on the configured host and port and serves plain TCP clients +// until ctx is cancelled, then shuts down gracefully (see serve). +func (s *Server) Run(ctx context.Context) error { + addr := net.JoinHostPort(s.conf.Host, strconv.Itoa(s.conf.Port)) + listen, err := net.Listen("tcp", addr) + if err != nil { + return fmt.Errorf("listen %s: %w", addr, err) + } + return s.serve(ctx, listen) +} + +// serve accepts connections on listen until ctx is cancelled or listen fails, +// then shuts down: it stops accepting, then closes the hub's clients +// gracefully, waiting up to the shutdown timeout (for ever if it is 0 or +// less). It returns nil after a shutdown caused by ctx, or the error that +// stopped it accepting, joined with any error from the hub's shutdown. +func (s *Server) serve(ctx context.Context, listen net.Listener) error { + // Clients outlive ctx until the hub has closed them, so the shutdown can + // tell them why; cancelling ctx itself would drop them without a word. + connCtx, cancelConns := context.WithCancel(context.WithoutCancel(ctx)) + defer cancelConns() + + stop := context.AfterFunc(ctx, func() { + listen.Close() //nolint:errcheck,gosec // Accept returns net.ErrClosed, which ends the loop + }) + defer stop() + + s.log.Info("listening", "addr", listen.Addr().String()) + acceptErr := s.accept(ctx, connCtx, listen) + if ctx.Err() != nil { + acceptErr = nil // cancelling ctx is how the server is stopped + } + listen.Close() //nolint:errcheck,gosec // closed already when ctx was cancelled + + s.log.Info("shutting down") + shutdownCtx := context.Background() + if s.conf.Timeouts.ShutdownTimeout > 0 { + var cancel context.CancelFunc + shutdownCtx, cancel = context.WithTimeout( + shutdownCtx, s.conf.Timeouts.ShutdownTimeout, + ) + defer cancel() + } + + errs := []error{acceptErr} + if err := s.hub.Shutdown(shutdownCtx); err != nil { + errs = append(errs, fmt.Errorf("websocket shutdown: %w", err)) + } + return errors.Join(errs...) +} + +// accept runs each connection on listen as a hub client, with connCtx, until +// ctx is cancelled (it returns ctx's error) or listen is closed some other way +// (it returns the error). Other accept errors are retried after a back off. +// Connections over the connection limit are turned away (see reject). +func (s *Server) accept( + ctx, connCtx context.Context, listen net.Listener, +) error { + limiter := s.connectionLimiter(ctx) + + var delay time.Duration + for { + newConn, err := listen.Accept() + if err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + if errors.Is(err, net.ErrClosed) { + return fmt.Errorf("accept: %w", err) + } + + // Back off and retry e.g. when out of file descriptors. + delay = min(max(2*delay, 5*time.Millisecond), time.Second) + s.log.Warn("tcp accept failed", "error", err, "retry_in", delay) + + select { + case <-time.After(delay): + continue + case <-ctx.Done(): + return ctx.Err() + } + } + delay = 0 + + if limiter != nil { + ip := remoteIP(newConn) + if allowed, _ := limiter.Allow(ip); !allowed { + s.log.Warn("tcp connection rate limited", "ip", ip) + go reject(newConn, "rate limited") + continue + } + } + + go s.hub.ServeTCP(connCtx, newConn) + } +} + +// connectionLimiter returns a limiter of new connections per client IP for +// websocket.connection_limit, or nil if it is off. Its idle entries are +// evicted after rate_limits.idle_ttl until ctx is cancelled. +func (s *Server) connectionLimiter(ctx context.Context) middleware.Limiter { + rule := s.conf.WebSocket.ConnectionLimit + if !rule.Enabled || rule.PerSecond <= 0 { + return nil + } + return middleware.NewKeyedLimiter( + ctx, float64(rule.PerSecond), max(rule.Burst, 1), s.conf.RateLimits.IdleTTL, + ) +} + +// reject tells a client why it is being turned away, with an error envelope +// as the hub would send, then closes its connection. +func reject(conn net.Conn, reason string) { + defer conn.Close() //nolint:errcheck // turned away; nothing to do if closing fails + + message, err := json.Marshal(websocket.Envelope{ + Type: websocket.TypeError, + Error: reason, + }) + if err != nil { + return + } + + conn.SetWriteDeadline(time.Now().Add(rejectTimeout)) //nolint:errcheck,gosec // best effort + conn.Write(append(message, '\n')) //nolint:errcheck,gosec // best effort +} + +// remoteIP returns the address of the connection's peer, or "" if it has no +// IP address. +func remoteIP(conn net.Conn) string { + addrPort, err := netip.ParseAddrPort(conn.RemoteAddr().String()) + if err != nil { + return "" + } + return addrPort.Addr().Unmap().String() +} diff --git a/internal/websocket/tcpserver/server_test.go b/internal/websocket/tcpserver/server_test.go new file mode 100644 index 0000000..be4dbfd --- /dev/null +++ b/internal/websocket/tcpserver/server_test.go @@ -0,0 +1,263 @@ +package tcpserver + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net" + "strings" + "syscall" + "testing" + "time" + + "git.avsdev.uk/Go/job-search/internal/config" + "git.avsdev.uk/Go/job-search/internal/websocket" +) + +var discardLog = slog.New(slog.NewTextHandler(io.Discard, nil)) + +// testConf is a server configuration for tests, changed by mutate if it is +// not nil. +func testConf(mutate func(*config.JSCServer)) *config.JSCServer { + c := &config.JSCServer{ + Host: "127.0.0.1", + Port: 0, + WebSocket: config.JSCServerWebSocket{WriteTimeout: time.Second}, + Timeouts: config.JSCServerTimeouts{ShutdownTimeout: time.Second}, + } + if mutate != nil { + mutate(c) + } + return c +} + +// newServer returns a server for conf with its own hub. +func newServer(conf *config.JSCServer) (*Server, *websocket.Hub) { + h := websocket.NewHub(conf.WebSocket, discardLog) + return New(conf, h, discardLog), h +} + +// startServer serves on a local port, returning the server's hub, its +// address, a function to cancel it and a channel with serve's result. +func startServer( + t *testing.T, conf *config.JSCServer, +) (*websocket.Hub, string, context.CancelFunc, chan error) { + t.Helper() + s, h := newServer(conf) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + done := make(chan error, 1) + go func() { done <- s.serve(ctx, ln) }() + return h, ln.Addr().String(), cancel, done +} + +// tcpClient is a plain TCP test client reading newline-delimited envelopes. +type tcpClient struct { + t *testing.T + conn net.Conn + r *bufio.Reader +} + +func dial(t *testing.T, addr string) *tcpClient { + t.Helper() + conn, err := net.Dial("tcp", addr) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { conn.Close() }) + return &tcpClient{t: t, conn: conn, r: bufio.NewReader(conn)} +} + +// read returns the next envelope, failing the test after a second. +func (c *tcpClient) read() websocket.Envelope { + c.t.Helper() + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + line, err := c.r.ReadString('\n') + if err != nil { + c.t.Fatalf("read: %v (got %q)", err, line) + } + var env websocket.Envelope + if err := json.Unmarshal([]byte(line), &env); err != nil { + c.t.Fatalf("bad envelope %q: %v", line, err) + } + return env +} + +// waitClosed checks the server closes the connection, skipping any lines +// sent first. A reset counts as closed. +func (c *tcpClient) waitClosed() { + c.t.Helper() + c.conn.SetReadDeadline(time.Now().Add(time.Second)) + if _, err := io.ReadAll(c.r); err != nil && !errors.Is(err, syscall.ECONNRESET) { + c.t.Fatalf("connection not closed: %v", err) + } +} + +// Clients get pushed messages; cancelling the server's context tells them it +// is shutting down, closes them, stops accepting and returns nil. +func TestServerShutdown(t *testing.T) { + h, addr, cancel, done := startServer(t, testConf(nil)) + c := dial(t, addr) + if e := c.read(); e.Type != websocket.TypeWelcome { + t.Fatalf("want welcome, got %+v", e) + } + for h.Count() != 1 { + time.Sleep(time.Millisecond) + } + h.Push("item.created", map[string]string{"id": "x"}) + if e := c.read(); e.Type != "item.created" { + t.Fatalf("want item.created, got %+v", e) + } + + cancel() + if e := c.read(); e.Type != websocket.TypeError || e.Error != "server shutting down" { + t.Fatalf("want shutdown error, got %+v", e) + } + c.waitClosed() + select { + case err := <-done: + if err != nil { + t.Fatalf("serve = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("serve did not return") + } + if conn, err := net.Dial("tcp", addr); err == nil { + conn.Close() + t.Error("listener still accepting") + } +} + +// Over websocket.connection_limit, a client is told it is rate limited and +// closed; other clients are unaffected. +func TestServerConnectionLimit(t *testing.T) { + conf := testConf(func(c *config.JSCServer) { + c.WebSocket.ConnectionLimit = config.JSCRateLimitRule{ + Enabled: true, PerSecond: 1, Burst: 1, + } + c.RateLimits.IdleTTL = time.Minute + }) + _, addr, cancel, _ := startServer(t, conf) + defer cancel() + + first := dial(t, addr) + if e := first.read(); e.Type != websocket.TypeWelcome { + t.Fatalf("first: want welcome, got %+v", e) + } + second := dial(t, addr) + if e := second.read(); e.Type != websocket.TypeError || e.Error != "rate limited" { + t.Fatalf("second: want rate limited, got %+v", e) + } + second.waitClosed() + + first.conn.Write([]byte(`{"type":"ping","id":"1"}` + "\n")) + if e := first.read(); e.Type != websocket.TypePong { + t.Fatalf("first: want pong, got %+v", e) + } +} + +// With the connection limit off (or no rate), there is no limiter. +func TestServerConnectionLimitOff(t *testing.T) { + for _, rule := range []config.JSCRateLimitRule{ + {Enabled: false, PerSecond: 1, Burst: 1}, + {Enabled: true, PerSecond: 0, Burst: 1}, + } { + s, _ := newServer(testConf(func(c *config.JSCServer) { + c.WebSocket.ConnectionLimit = rule + })) + if l := s.connectionLimiter(context.Background()); l != nil { + t.Errorf("rule %+v: got a limiter", rule) + } + } +} + +func TestServerRun(t *testing.T) { + s, _ := newServer(testConf(nil)) + ctx, cancel := context.WithCancel(context.Background()) + time.AfterFunc(20*time.Millisecond, cancel) + if err := s.Run(ctx); err != nil { + t.Fatalf("Run = %v", err) + } +} + +func TestServerRunListenError(t *testing.T) { + s, _ := newServer(testConf(func(c *config.JSCServer) { c.Port = -1 })) + if err := s.Run(context.Background()); err == nil || !strings.Contains(err.Error(), "listen") { + t.Fatalf("Run = %v; want a listen error", err) + } +} + +// flakyListener fails Accept with a temporary error a number of times, then +// as if closed. +type flakyListener struct { + net.Listener + fails int +} + +func (l *flakyListener) Accept() (net.Conn, error) { + if l.fails > 0 { + l.fails-- + return nil, errors.New("too many open files") + } + return nil, net.ErrClosed +} + +// Accept errors other than a closed listener are retried after a back off. +func TestServerAcceptRetries(t *testing.T) { + s, _ := newServer(testConf(nil)) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + fl := &flakyListener{Listener: ln, fails: 3} + err = s.accept(context.Background(), context.Background(), fl) + if !errors.Is(err, net.ErrClosed) || fl.fails != 0 { + t.Fatalf("accept = %v with %d failures left; want net.ErrClosed", err, fl.fails) + } +} + +// Cancelling ctx stops accept while it is backing off. +func TestServerAcceptCancelledInBackOff(t *testing.T) { + s, _ := newServer(testConf(nil)) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + ctx, cancel := context.WithCancel(context.Background()) + fl := &flakyListener{Listener: ln, fails: 1 << 30} + time.AfterFunc(20*time.Millisecond, cancel) + if err := s.accept(ctx, ctx, fl); !errors.Is(err, context.Canceled) { + t.Fatalf("accept = %v; want context.Canceled", err) + } +} + +// A listener closed from elsewhere is an error, after a graceful shutdown. +func TestServerListenerClosed(t *testing.T) { + s, _ := newServer(testConf(nil)) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + ln.Close() + if err := s.serve(context.Background(), ln); !errors.Is(err, net.ErrClosed) { + t.Fatalf("serve = %v; want net.ErrClosed", err) + } +} + +func TestRemoteIP(t *testing.T) { + a, b := net.Pipe() + defer a.Close() + defer b.Close() + if ip := remoteIP(a); ip != "" { + t.Errorf("remoteIP(pipe) = %q; want empty", ip) + } +}