Go で WebSocket を簡単に扱えるパッケージを作った

趣味で Go 言語で WebSocket サーバーを書いてるんですが、そこで使ってる WebSocket コネクションのハンドリング部分をパッケージとして切り出したら便利なんじゃないかと思って切り出してみました。ベースは github.com/gorilla/websocket です。

インストール方法

$ go get github.com/shiwano/websocket-conn

使い方

func Connect(ctx context.Context, settings Settings, url string, requestHeader http.Header) (*Conn *http.Response, error)
func UpgradeFromHTTP(ctx context.Context, settings Settings, w http.ResponseWriter, r *http.Request) (*Conn, error)

type Conn struct {
	Stream() <-chan Data
	Err() error
	SendBinaryMessage(data []byte) error
	SendTextMessage(text string) error
}

サンプルコード

サーバー側:

package main

import (
	"context"
	"net/http"
	"github.com/shiwano/websocket-conn"
)

func main() {
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		c, err := conn.UpgradeFromHTTP(ctx, conn.DefaultSettings(), w, r)
		if err != nil {
			w.Write([]byte("Error"))
			return
		}
		d := <-c.Stream()
		c.SendTextMessage(d.Message.Text() + " World")
		cancel()
	})
	http.ListenAndServe(":5000", nil)
}

クライアント側:

package main

import (
	"context"
	"log"
	"github.com/shiwano/websocket-conn"
)

func main() {
	c, _, err := conn.Connect(context.Background(), conn.DefaultSettings(), "ws://localhost:5000", nil)
	if err != nil {
		log.Fatal(err)
	}
	c.SendTextMessage("Hello")
	d := <-c.Stream()
	log.Println(d.Message.Text()) // Hello World
	for d := range c.Stream() {
		if d.EOS {
			log.Println("Closed: ", c.Err())
		}
	}
}

上記のコードをそれぞれ実行したら、クライアント側で Hello World と出力されるはずです。こんな感じでけっこうお手軽に WebSocket サーバーとそのクライアントを作れるので、よかったら使ってみてください(´・ω・`)