diff --git a/streaming_ws.go b/streaming_ws.go new file mode 100644 index 0000000..915465f --- /dev/null +++ b/streaming_ws.go @@ -0,0 +1,178 @@ +package mastodon + +import ( + "context" + "encoding/json" + "net/url" + "path" + + "github.com/gorilla/websocket" +) + +// WSClient is a WebSocket client. +type WSClient struct { + websocket.Dialer + client *Client +} + +// NewWSClient return WebSocket client. +func (c *Client) NewWSClient() *WSClient { return &WSClient{client: c} } + +// Stream is a struct of data that flows in streaming. +type Stream struct { + Event string `json:"event"` + Payload interface{} `json:"payload"` +} + +// StreamingWSPublic return channel to read events on public using WebSocket. +func (c *WSClient) StreamingWSPublic(ctx context.Context) (chan Event, error) { + return c.streamingWS(ctx, "public", "") +} + +// StreamingWSPublicLocal return channel to read events on public local using WebSocket. +func (c *WSClient) StreamingWSPublicLocal(ctx context.Context) (chan Event, error) { + return c.streamingWS(ctx, "public:local", "") +} + +// StreamingWSUser return channel to read events on home using WebSocket. +func (c *WSClient) StreamingWSUser(ctx context.Context) (chan Event, error) { + return c.streamingWS(ctx, "user", "") +} + +// StreamingWSHashtag return channel to read events on tagged timeline using WebSocket. +func (c *WSClient) StreamingWSHashtag(ctx context.Context, tag string) (chan Event, error) { + return c.streamingWS(ctx, "hashtag", tag) +} + +// StreamingWSHashtagLocal return channel to read events on tagged local timeline using WebSocket. +func (c *WSClient) StreamingWSHashtagLocal(ctx context.Context, tag string) (chan Event, error) { + return c.streamingWS(ctx, "hashtag:local", tag) +} + +func (c *WSClient) streamingWS(ctx context.Context, stream, tag string) (chan Event, error) { + params := url.Values{} + params.Set("access_token", c.client.config.AccessToken) + params.Set("stream", stream) + if tag != "" { + params.Set("tag", tag) + } + + u, err := changeWebSocketScheme(c.client.config.Server) + if err != nil { + return nil, err + } + u.Path = path.Join(u.Path, "/api/v1/streaming") + u.RawQuery = params.Encode() + + q := make(chan Event) + go func() { + for { + err := c.handleWS(ctx, u.String(), q) + if err != nil { + return + } + } + }() + + return q, nil +} + +func (c *WSClient) handleWS(ctx context.Context, rawurl string, q chan Event) error { + conn, err := c.dialRedirect(rawurl) + if err != nil { + q <- &ErrorEvent{err: err} + + // End. + return err + } + defer conn.Close() + + for { + select { + case <-ctx.Done(): + q <- &ErrorEvent{err: ctx.Err()} + + // End. + return ctx.Err() + default: + } + + var s Stream + err := conn.ReadJSON(&s) + if err != nil { + q <- &ErrorEvent{err: err} + + // Reconnect. + break + } + + err = nil + switch s.Event { + case "update": + var status Status + err = json.Unmarshal([]byte(s.Payload.(string)), &status) + if err == nil { + q <- &UpdateEvent{Status: &status} + } + case "notification": + var notification Notification + err = json.Unmarshal([]byte(s.Payload.(string)), ¬ification) + if err == nil { + q <- &NotificationEvent{Notification: ¬ification} + } + case "delete": + q <- &DeleteEvent{ID: int64(s.Payload.(float64))} + } + if err != nil { + q <- &ErrorEvent{err} + } + } + + return nil +} + +func (c *WSClient) dialRedirect(rawurl string) (conn *websocket.Conn, err error) { + for { + conn, rawurl, err = c.dial(rawurl) + if err != nil { + return nil, err + } else if conn != nil { + return conn, nil + } + } +} + +func (c *WSClient) dial(rawurl string) (*websocket.Conn, string, error) { + conn, resp, err := c.Dial(rawurl, nil) + if err != nil && err != websocket.ErrBadHandshake { + return nil, "", err + } + defer resp.Body.Close() + + if loc := resp.Header.Get("Location"); loc != "" { + u, err := changeWebSocketScheme(loc) + if err != nil { + return nil, "", err + } + + return nil, u.String(), nil + } + + return conn, "", err +} + +func changeWebSocketScheme(rawurl string) (*url.URL, error) { + u, err := url.Parse(rawurl) + if err != nil { + return nil, err + } + + switch u.Scheme { + case "http": + u.Scheme = "ws" + case "https": + u.Scheme = "wss" + } + + return u, nil +} diff --git a/streaming_ws_test.go b/streaming_ws_test.go new file mode 100644 index 0000000..5456a5d --- /dev/null +++ b/streaming_ws_test.go @@ -0,0 +1,298 @@ +package mastodon + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +func TestStreamingWSPublic(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(wsMock)) + defer ts.Close() + + client := NewClient(&Config{Server: ts.URL}).NewWSClient() + ctx, cancel := context.WithCancel(context.Background()) + q, err := client.StreamingWSPublic(ctx) + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + + wsTest(t, q, cancel) +} + +func TestStreamingWSPublicLocal(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(wsMock)) + defer ts.Close() + + client := NewClient(&Config{Server: ts.URL}).NewWSClient() + ctx, cancel := context.WithCancel(context.Background()) + q, err := client.StreamingWSPublicLocal(ctx) + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + + wsTest(t, q, cancel) +} + +func TestStreamingWSUser(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(wsMock)) + defer ts.Close() + + client := NewClient(&Config{Server: ts.URL}).NewWSClient() + ctx, cancel := context.WithCancel(context.Background()) + q, err := client.StreamingWSUser(ctx) + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + + wsTest(t, q, cancel) +} + +func TestStreamingWSHashtag(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(wsMock)) + defer ts.Close() + + client := NewClient(&Config{Server: ts.URL}).NewWSClient() + ctx, cancel := context.WithCancel(context.Background()) + q, err := client.StreamingWSHashtag(ctx, "zzz") + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + + wsTest(t, q, cancel) +} + +func TestStreamingWSHashtagLocal(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(wsMock)) + defer ts.Close() + + client := NewClient(&Config{Server: ts.URL}).NewWSClient() + ctx, cancel := context.WithCancel(context.Background()) + q, err := client.StreamingWSHashtagLocal(ctx, "zzz") + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + + wsTest(t, q, cancel) +} + +func wsMock(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/v1/streaming" { + http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound) + return + } + + u := websocket.Upgrader{} + conn, err := u.Upgrade(w, r, nil) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + defer conn.Close() + + err = conn.WriteMessage(websocket.TextMessage, + []byte(`{"event":"update","payload":"{\"content\":\"foo\"}"}`)) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + + err = conn.WriteMessage(websocket.TextMessage, + []byte(`{"event":"notification","payload":"{\"id\":123}"}`)) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + + err = conn.WriteMessage(websocket.TextMessage, + []byte(`{"event":"delete","payload":1234567}`)) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + + err = conn.WriteMessage(websocket.TextMessage, + []byte(`{"event":"update","payload":""}`)) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + + time.Sleep(10 * time.Second) +} + +func wsTest(t *testing.T, q chan Event, cancel func()) { + time.AfterFunc(time.Second, func() { + cancel() + close(q) + }) + events := []Event{} + for e := range q { + events = append(events, e) + } + if len(events) != 4 { + t.Fatalf("result should be two: %d", len(events)) + } + if events[0].(*UpdateEvent).Status.Content != "foo" { + t.Fatalf("want %q but %q", "foo", events[0].(*UpdateEvent).Status.Content) + } + if events[1].(*NotificationEvent).Notification.ID != 123 { + t.Fatalf("want %d but %d", 123, events[1].(*NotificationEvent).Notification.ID) + } + if events[2].(*DeleteEvent).ID != 1234567 { + t.Fatalf("want %d but %d", 1234567, events[2].(*DeleteEvent).ID) + } + if errorEvent, ok := events[3].(*ErrorEvent); !ok { + t.Fatalf("should be fail: %v", errorEvent.err) + } +} + +func TestStreamingWS(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(wsMock)) + defer ts.Close() + + client := NewClient(&Config{Server: ":"}).NewWSClient() + _, err := client.StreamingWSPublicLocal(context.Background()) + if err == nil { + t.Fatalf("should be fail: %v", err) + } + + client = NewClient(&Config{Server: ts.URL}).NewWSClient() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + q, err := client.StreamingWSPublicLocal(ctx) + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + go func() { + e := <-q + if errorEvent, ok := e.(*ErrorEvent); !ok { + t.Fatalf("should be fail: %v", errorEvent.err) + } + }() +} + +func TestHandleWS(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + u := websocket.Upgrader{} + conn, err := u.Upgrade(w, r, nil) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + defer conn.Close() + + err = conn.WriteMessage(websocket.TextMessage, + []byte(``)) + if err != nil { + http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError) + return + } + + time.Sleep(10 * time.Second) + })) + defer ts.Close() + + q := make(chan Event) + client := NewClient(&Config{}).NewWSClient() + + go func() { + e := <-q + if errorEvent, ok := e.(*ErrorEvent); !ok { + t.Fatalf("should be fail: %v", errorEvent.err) + } + }() + err := client.handleWS(context.Background(), ":", q) + if err == nil { + t.Fatalf("should be fail: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + go func() { + e := <-q + if errorEvent, ok := e.(*ErrorEvent); !ok { + t.Fatalf("should be fail: %v", errorEvent.err) + } + }() + err = client.handleWS(ctx, "ws://"+ts.Listener.Addr().String(), q) + if err == nil { + t.Fatalf("should be fail: %v", err) + } + + go func() { + e := <-q + if errorEvent, ok := e.(*ErrorEvent); !ok { + t.Fatalf("should be fail: %v", errorEvent.err) + } + }() + client.handleWS(context.Background(), "ws://"+ts.Listener.Addr().String(), q) +} + +func TestDialRedirect(t *testing.T) { + client := NewClient(&Config{}).NewWSClient() + _, err := client.dialRedirect(":") + if err == nil { + t.Fatalf("should be fail: %v", err) + } +} + +func TestDial(t *testing.T) { + canErr := true + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if canErr { + canErr = false + http.Redirect(w, r, ":", http.StatusMovedPermanently) + return + } + + http.Redirect(w, r, "http://www.example.com/", http.StatusMovedPermanently) + })) + defer ts.Close() + + client := NewClient(&Config{}).NewWSClient() + _, _, err := client.dial(":") + if err == nil { + t.Fatalf("should be fail: %v", err) + } + + _, rawurl, err := client.dial("ws://" + ts.Listener.Addr().String()) + if err == nil { + t.Fatalf("should be fail: %v", err) + } + + _, rawurl, err = client.dial("ws://" + ts.Listener.Addr().String()) + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + if rawurl != "ws://www.example.com/" { + t.Fatalf("want %q but %q", "ws://www.example.com/", rawurl) + } +} + +func TestChangeWebSocketScheme(t *testing.T) { + _, err := changeWebSocketScheme(":") + if err == nil { + t.Fatalf("should be fail: %v", err) + } + + u, err := changeWebSocketScheme("http://example.com/") + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + if u.Scheme != "ws" { + t.Fatalf("want %q but %q", "ws", u.Scheme) + } + + u, err = changeWebSocketScheme("https://example.com/") + if err != nil { + t.Fatalf("should not be fail: %v", err) + } + if u.Scheme != "wss" { + t.Fatalf("want %q but %q", "wss", u.Scheme) + } +}