jsonrpc on websocket for golang

client.go

// client.go
package main

import (
	"fmt"
	"log"
	"net/rpc"
	"net/rpc/jsonrpc"

	"github.com/gorilla/websocket"
)

// 参数结构体
type Args struct {
	A int
	B int
}
type WebSocketConn struct {
	conn *websocket.Conn
}

func (wsc *WebSocketConn) Read(p []byte) (n int, err error) {
	_, msg, err := wsc.conn.ReadMessage()
	if err != nil {
		return 0, err
	}
	copy(p, msg)
	return len(msg), nil
}

func (wsc *WebSocketConn) Write(p []byte) (n int, err error) {
	err = wsc.conn.WriteMessage(websocket.TextMessage, p)
	if err != nil {
		return 0, err
	}
	return len(p), nil
}

func (wsc *WebSocketConn) Close() error {
	return wsc.conn.Close()
}

func main() {

	conn, _, err := websocket.DefaultDialer.Dial("ws://localhost:7000/ws", nil)
	if err != nil {
		log.Fatal("Error connecting to server:", err)
	}
	defer conn.Close()

	wsc := &WebSocketConn{conn: conn}

	args := &Args{A: 3, B: 4}
	var reply int
	client := rpc.NewClientWithCodec(jsonrpc.NewClientCodec(wsc))

	err = client.Call("Arith.Multiply", args, &reply)
	if err != nil {
		log.Fatal("Multiply error:", err)
	}
	fmt.Printf("Multiply: %d*%d=%d\n", args.A, args.B, reply)
}

server.go

package main

import (
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"

	"github.com/gorilla/websocket"
)

type Args struct {
	A int
	B int
}

type Arith int

func (t *Arith) Multiply(args *Args, reply *int) error {
	*reply = args.A * args.B
	return nil
}

type WebSocketConn struct {
	conn *websocket.Conn
}

func (wsc *WebSocketConn) Read(p []byte) (n int, err error) {
	_, msg, err := wsc.conn.ReadMessage()
	if err != nil {
		return 0, err
	}
	copy(p, msg)
	return len(msg), nil
}

func (wsc *WebSocketConn) Write(p []byte) (n int, err error) {
	err = wsc.conn.WriteMessage(websocket.TextMessage, p)
	if err != nil {
		return 0, err
	}
	return len(p), nil
}

func (wsc *WebSocketConn) Close() error {
	return wsc.conn.Close()
}

var (
	upgrader = websocket.Upgrader{}
	rpcServer = rpc.NewServer()
	
)

// WebSocket 处理函数
func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		http.Error(w, "Failed to upgrade connection", http.StatusBadRequest)
		return
	}

	go rpcServer.ServeCodec(jsonrpc.NewServerCodec(&WebSocketConn{conn: conn}))
}

func main() {
	rpcServer.Register(new(Arith))
	http.HandleFunc("/ws", wsHandler)
	http.ListenAndServe(":7000", nil)
}

ref: https://gist.github.com/mowings/dad4917b6fc21e6e5ffe2da3b16e7ef9