1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
|
package main
import (
"context"
"fmt"
"io"
"net/http"
"time"
"golang.org/x/time/rate"
"github.com/coder/websocket"
)
// echoServer is the WebSocket echo server implementation.
// It ensures the client speaks the echo subprotocol and
// only allows one message every 100ms with a 10 message burst.
type echoServer struct {
// logf controls where logs are sent.
logf func(f string, v ...interface{})
}
func (s echoServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
c, err := websocket.Accept(w, r, &websocket.AcceptOptions{
Subprotocols: []string{"echo"},
})
if err != nil {
s.logf("%v", err)
return
}
defer c.CloseNow()
if c.Subprotocol() != "echo" {
c.Close(websocket.StatusPolicyViolation, "client must speak the echo subprotocol")
return
}
l := rate.NewLimiter(rate.Every(time.Millisecond*100), 10)
for {
err = echo(r.Context(), c, l)
if websocket.CloseStatus(err) == websocket.StatusNormalClosure {
return
}
if err != nil {
s.logf("failed to echo with %v: %v", r.RemoteAddr, err)
return
}
}
}
// echo reads from the WebSocket connection and then writes
// the received message back to it.
// The entire function has 10s to complete.
func echo(ctx context.Context, c *websocket.Conn, l *rate.Limiter) error {
ctx, cancel := context.WithTimeout(ctx, time.Second*10)
defer cancel()
err := l.Wait(ctx)
if err != nil {
return err
}
typ, r, err := c.Reader(ctx)
if err != nil {
return err
}
w, err := c.Writer(ctx, typ)
if err != nil {
return err
}
_, err = io.Copy(w, r)
if err != nil {
return fmt.Errorf("failed to io.Copy: %w", err)
}
err = w.Close()
return err
}
|
