Added websocket hub for handling websocket clients, with a wsserve command
which runs a plain TCP server for receiving websocket hub clients, with server config, CSRF and rate limit middleware. - internal/websocket: Hub pushing messages to WebSocket clients (ServeHTTP, example.com/ws upgrade with optional CSRF check) and plain TCP clients (ServeTCP, newline-delimited JSON envelopes); ping/pong keepalive, per-client message rate limit, slow consumer handling, graceful Shutdown, debug logging of messages sent and received - internal/websocket/tcpserver: standalone TCP server for a hub, with accept back-off, a per-IP connection limit (server.websocket.connection_limit) and graceful shutdown - internal/cli/wsserve: `wsserve` command running the TCP server, with --listen/--host/--port overrides - internal/middleware: keyed token bucket rate limiter, real client IP from trusted proxies - internal/httpx: RFC 9457 problem responses, request IDs in context - internal/config: server section (host, port, shutdown timeout, rate limiter idle TTL, websocket); util.ByteSize settings (e.g. "32KiB"); lists of sections as single keys; empty lists stored as nil; isExplicitlySet for printLoadedConfig; JobSearchConfigSource renamed to JSCSourcesSource; TestSections checks every section's methods - internal/util: ByteSize, parsed and formatted with go-humanize
This commit is contained in:
bwilliams committed
1 parent
803a8929e6
commit
c9f42b62ad
42 files changed
+3619
-57
No files matched your search
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -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=
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 |
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
|
||||
@@ -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) }
|
||||
@@ -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
|
||||
}
|
||||
|
||||
+122
-15
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+186
-16
@@ -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
|
||||
|
||||
@@ -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" {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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{}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user