From c7255e69e858573604b449f65c205a430fb05e1a Mon Sep 17 00:00:00 2001
From: 178inaba <178inaba@users.noreply.github.com>
Date: Sun, 23 Apr 2017 15:42:57 +0900
Subject: [PATCH] Add WebSocket implementation of streaming API

---
 mastodon.go          |   3 +
 streaming_ws.go      | 171 +++++++++++++++++++++++++++++++++++++++++++
 streaming_ws_test.go |   1 +
 3 files changed, 175 insertions(+)
 create mode 100644 streaming_ws.go
 create mode 100644 streaming_ws_test.go

diff --git a/mastodon.go b/mastodon.go
index 16ccb08..594d16a 100644
--- a/mastodon.go
+++ b/mastodon.go
@@ -13,6 +13,8 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+
+	"github.com/gorilla/websocket"
 )
 
 // Config is a setting for access mastodon APIs.
@@ -26,6 +28,7 @@ type Config struct {
 // Client is a API client for mastodon.
 type Client struct {
 	http.Client
+	websocket.Dialer
 	config   *Config
 	interval time.Duration
 }
diff --git a/streaming_ws.go b/streaming_ws.go
new file mode 100644
index 0000000..20dd0a1
--- /dev/null
+++ b/streaming_ws.go
@@ -0,0 +1,171 @@
+package mastodon
+
+import (
+	"context"
+	"encoding/json"
+	"net/url"
+	"path"
+
+	"github.com/gorilla/websocket"
+)
+
+// 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 *Client) 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 *Client) 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 *Client) StreamingWSUser(ctx context.Context, user string) (chan Event, error) {
+	return c.streamingWS(ctx, "user", "")
+}
+
+// StreamingWSHashtag return channel to read events on tagged timeline using WebSocket.
+func (c *Client) 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 *Client) StreamingWSHashtagLocal(ctx context.Context, tag string) (chan Event, error) {
+	return c.streamingWS(ctx, "hashtag:local", tag)
+}
+
+func (c *Client) streamingWS(ctx context.Context, stream, tag string) (chan Event, error) {
+	params := url.Values{}
+	params.Set("access_token", c.config.AccessToken)
+	params.Set("stream", stream)
+	if tag != "" {
+		params.Set("tag", tag)
+	}
+
+	u, err := changeWebSocketScheme(c.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 *Client) 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)), &notification)
+			if err == nil {
+				q <- &NotificationEvent{Notification: &notification}
+			}
+		case "delete":
+			if err == nil {
+				q <- &DeleteEvent{ID: int64(s.Payload.(float64))}
+			}
+		}
+		if err != nil {
+			q <- &ErrorEvent{err}
+		}
+	}
+
+	return nil
+}
+
+func (c *Client) 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 *Client) 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..6cd653a
--- /dev/null
+++ b/streaming_ws_test.go
@@ -0,0 +1 @@
+package mastodon