【Golang】WebSocketを使用した簡単なチャットの実装

# GoでのWebSocket実装

通知機能の実装やチャットアプリなどで使用される、WebSocketについて学習してみたいと思います。
今回の目標は、JWTを使って認証後、WebSocket接続を確立し、クライアント側とサーバー側で双方向にチャットができるようにすることです。
クライアントはHTMLを用意し、サーバー側はターミナルから入力することにします。

## WebSocketの仕組み（HTTPとの主な違い）
WebSocketとは、クライアント（ブラウザなど）とサーバー間で双方向通信を可能にするプロトコル。
通常のHTTP通信とは異なり、一度接続が確立されると、その接続を維持しながらリアルタイムにデータの送受信ができる。
### 主な違い
| 項目 | HTTP | WebSocket |
|:-:|:-:|:-:|
| 通信方式 | リクエスト・レスポンス型（クライアントがリクエストを送信し、サーバーが応答する） | 双方向通信（どちらからでもデータを送信できる |
| 接続の維持 | 通信ごとに新しい接続を確立 | 一度確立すると、維持される |
| リアルタイム性 | リクエストしない限りデータを取得できない（ポーリングが必要） | サーバー側からのプッシュが可能 |
| プロトコル | HTTP/1.1, HTTP/2, HTTP/3 | ws://またはwss:// |
| 使用例 | REST API, Webページのデータ取得 | チャットアプリ、オンラインゲーム、リアルタイム通知 |

### WebSocketの利点
- 通信のオーバーヘッドが小さい（HTTPのヘッダー情報を何度も送る必要がない）
- サーバープッシュが可能（サーバーからクライアントにデータを即座に送信できる）

### JWTを用いたWebSocketの仕組み
WebSocket自体には標準的な認証の仕組みがないため、トークン（JWT）を使って認証を行う。
JWT（JSON Web Token）は、ユーザーの情報を含む暗号化されたトークンで、認証と認可に使用される。

### 認証の流れ
1. クライアントがJWTを取得
    - クライアントは、ログイン時などにサーバーからJWTを取得（HTTPの/loginエンドポイントなど）
    - 取得したJWTは、後のWebSocket通信で使用される
2. WebSocket接続時にJWTを送信
    - JWTの送信方法は、クエリパラメータで送信する方法と、ヘッダー（Sec-WebSocket-Protocol）を使用する方法がある
3. サーバー側でJWTを検証
    - サーバーはJWTの署名を検証し、トークンが有効かチェックする
    - トークンが有効なら、ユーザー情報を取得し、接続を許可する

※JWTには有効期限（expフレーム）があり、期限が切れた場合は再ログイン、または、リフレッシュトークンが必要

### お絵描き
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3771904/409916a3-3c8e-c0d1-a21c-83f8f00cb503.png)


## 実装例
### サーバー側（Go）
使用するライブラリのインストール
```sh
go get github.com/gorilla/websocket
go get github.com/golang-jwt/jwt/v5
```
```main.go
package main

import (
	"bufio"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/gorilla/websocket"
)

var jwtSecret = []byte("secret_key")

// JWT を検証
// JWTは３つの部分で構成される文字列
// header.payload.signature
// header：アルゴリズムやトークンのタイプを定義
// payload：ユーザー情報や有効期限などのデータ
// signature：トークンが改ざんされていないか確認するためのもの
func validateJWT(tokenString string) (*jwt.Token, error) {
	// jwt.Parse(tokenString, ...) でトークンを解析
	// tokenStringを分解し、header.payload.signatureの３つを取り出す
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		// 署名アルゴリズムの確認（HMACを使う前提でチェックしてる）
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method")
		}
		// JWTの署名部分を検証し、改ざんされていないかチェック
		// jwtSecret（サーバー側の秘密鍵）を使って、クライアントが送信した署名とサーバーで生成した署名を比較
		return jwtSecret, nil
	})

	// 署名が無効だったり、期限切れのトークンは拒否
	if err != nil || !token.Valid {
		return nil, fmt.Errorf("invalid token")
	}
	return token, nil
}

// WebSocket 接続用のアップグレーダー
// HTTP→WebSocketへの切り替えをアップグレードという
// HTTP接続をWebSocket接続に変更する
var upgrader = websocket.Upgrader{
	// どのオリジンからのWebSocket接続も許可する
	// ※gorilla/websocketでは、デフォルトではクロスオリジンからの接続をブロックするようになっている
	CheckOrigin: func(r *http.Request) bool { return true },
}

// WebSocket ハンドラー
func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	// JWT の検証
	// URLからtokenを取得
	tokenString := r.URL.Query().Get("token")
	if tokenString == "" {
		http.Error(w, "Missing token", http.StatusUnauthorized)
		return
	}

	_, err := validateJWT(tokenString)
	if err != nil {
		http.Error(w, "Invalid token", http.StatusUnauthorized)
		return
	}

	// WebSocket 接続
	// UpgradeはHTTPリクエストを解析し、101 Switching Protocols（プロトコルの変更）のレスポンスをクライアントに返す
	// 戻り値のconnには、WebSocket接続のインスタンスが格納され、双方向通信が可能となる
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("WebSocket upgrade error:", err)
		return
	}
	defer conn.Close()

	log.Println("Client connected")

	// ターミナル入力を監視するゴルーチン
	go func() {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			message := scanner.Text()
			// クライアントにメッセージを送信
			if err := conn.WriteMessage(websocket.TextMessage, []byte("[Server]: "+message)); err != nil {
				log.Println("Write error:", err)
				return
			}
		}
	}()

	// クライアントからのメッセージを受信
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			log.Println("Read error:", err)
			break
		}
		fmt.Println("[Client]:", string(msg))
	}
}

func generateJWT() (string, error) {
	// jwt.NewWithClaimsで新しいJWTを作成
	// SigningMethodHS256（HMAC-SHA256）で署名
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
		// 有効期限を1時間後に設定
		"exp": time.Now().Add(time.Hour * 1).Unix(),
	})
	// SignedString(jwtSecret)で署名し、最終的なJWTを生成
	return token.SignedString(jwtSecret)
}

func main() {
	// JWT を生成
	token, err := generateJWT()
	if err != nil {
		log.Fatal("Failed to generate JWT:", err)
	}
	// ここで表示したtokenをクライアント側にコピペする
	fmt.Println("Use this token in the frontend:", token)

	// WebSocket サーバー
	http.HandleFunc("/ws", handleWebSocket)
	log.Println("Server started on :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

```

### クライアント側（JavaScript）
```html:index.html
<!DOCTYPE html>
<html lang="ja">
    <head>
        <meta charset="UTF-8" />
        <meta name="viewport" content="width=device-width, initial-scale=1.0" />
        <title>WebSocket Chat</title>
    </head>
    <body>
        <h2>WebSocket Chat</h2>
        <input type="text" id="message" placeholder="Enter message" />
        <button onclick="sendMessage()">Send</button>
        <div id="chat"></div>

        <script>
            const token = prompt("Enter JWT Token:");
            const ws = new WebSocket(`ws://localhost:8080/ws?token=${token}`);

            ws.onopen = () => {
                console.log("Connected to WebSocket server.");
                document.getElementById("chat").innerHTML += "<p>Connected.</p>";
            };

            ws.onmessage = (event) => {
                document.getElementById("chat").innerHTML += `<p>${event.data}</p>`;
            };

            ws.onerror = (error) => {
                console.log("WebSocket Error:", error);
            };

            ws.onclose = () => {
                console.log("WebSocket connection closed.");
                document.getElementById("chat").innerHTML += "<p>Disconnected.</p>";
            };

            function sendMessage() {
                const message = document.getElementById("message").value;
                ws.send(message);
                document.getElementById("chat").innerHTML += `<p>You: ${message}</p>`;
                document.getElementById("message").value = "";
            }
        </script>
    </body>
</html>

```

### 使い方
サーバー起動時にターミナルに出力されたトークンをブラウザ側で入力すると、WebSocket接続ができる。
クライアントからサーバーにメッセージを送るときは、画面から入力し、Sendボタン押下。
サーバーからクライアントにメッセージを送るときは、ターミナルから入力し、Enterキー押下。

クライアントがJWTを受け取るのは、本来ならログインに成功した場合だが、今回は手動でコピペするようにした。

## 最後に
簡単ですが、WebSocketを使用したチャットを実装してみました。
HTTP以外の通信方法を使えるようになってワクワクしました！
引き続き学習に励みたいと思います！