一个WebSocket广播功能实现

一个WebSocket广播功能实现

ws.go

package main

import (
	"encoding/json"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/robfig/cron"
	"io/ioutil"
	_ "io/ioutil"
	"log"
	"net/http"
	"os"
)

type AppConfig struct {
	Data   string `json:"data"`
	Cron   string `json:"cron"`
	Url    string `json:"url"`
	Enable bool   `json:"enable"`
}

// ClientManager is a websocket manager
type ClientManager struct {
	Clients    map[string]*Client
	Broadcast  chan []byte
	Register   chan *Client
	Unregister chan *Client
}

// Client is a websocket client
type Client struct {
	ID     string
	Socket *websocket.Conn
	Send   chan []byte
}

// Message is return msg
type Message struct {
	Sender    string `json:"sender,omitempty"`
	Recipient string `json:"recipient,omitempty"`
	Content   string `json:"content,omitempty"`
}

// Manager define a ws server manager
var Manager = ClientManager{
	Broadcast:  make(chan []byte),
	Register:   make(chan *Client),
	Unregister: make(chan *Client),
	Clients:    make(map[string]*Client),
}
var AppConfigData AppConfig

var broadcastMsg []byte

var task *cron.Cron

func init() {
	log.Println("初始化WebSocket配置文件")
	// 读取 JSON 文件
	data, err := os.ReadFile("config.json")
	if err != nil {
		log.Println("初始化WebSocket配置文件失败...", err)
		panic(err)
	}

	// 解码 JSON 内容
	err = json.Unmarshal(data, &AppConfigData)
	if err != nil {
		panic(err)
	}

}

// Start is  项目运行前, 协程开启start -> go Manager.Start()
func (manager *ClientManager) Start() {
	log.Println("首页地址 : https://localhost:8080/")
	log.Println("在线测试地址：http://wstool.jackxiang.com/")
	createTask()
	for {
		log.Println("<---管道通信,等待连接中......--->")

		select {
		case conn := <-Manager.Register:
			log.Printf("新用户加入:%v", conn.ID)
			Manager.Clients[conn.ID] = conn
			jsonMessage, _ := json.Marshal(&Message{Content: "Successful connection to socket service"})
			conn.Send <- jsonMessage

		case conn := <-Manager.Unregister:
			log.Printf("用户离开:%v", conn.ID)
			if c, ok := Manager.Clients[conn.ID]; ok {
				c.Socket.Close()
				delete(Manager.Clients, conn.ID)
				if len(Manager.Clients) == 0 {
					log.Printf("所有的客户端断开连接")
				}
			}
		case message := <-Manager.Broadcast:
			//广播
			for _, conn := range Manager.Clients {
				conn.Send <- message
			}
		}
	}
}

func createUUID() string {
	return uuid.New().String()
}
func (c *Client) Read() {
	defer func() {
		//Manager.Unregister <- c
		c.Socket.Close()
	}()
	for {
		c.Socket.PongHandler()
		_, message, err := c.Socket.ReadMessage()
		if err != nil {
			Manager.Unregister <- c
			break
		}
		log.Printf("读取到客户端的信息:%s", string(message))
		//广播
		Manager.Broadcast <- message
		broadcastMsg = message
	}
}

func (c *Client) Write() {
	defer func() {
		c.Socket.Close()
	}()
	for {
		select {
		case message, ok := <-c.Send:
			if !ok {
				c.Socket.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			log.Printf("发送到到客户端:[%s]的信息:%s", c.ID, string(message))
			err := c.Socket.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				log.Printf("发送到到客户端:[%s]的信息:%s 失败", c.ID, string(message))
			}
		}
	}
}

func createTask() {
	task := cron.New()
	task.AddFunc(AppConfigData.Cron, func() {
		RunTimer(broadcastMsg)
	})
	task.Start()
}

func RunTimer(message []byte) {
	if AppConfigData.Enable {
		count := len(Manager.Clients)
		log.Printf("客户端数量：%d", count)
		for _, conn := range Manager.Clients {
			if message == nil {
				conn.Send <- []byte(AppConfigData.Data)
			} else {
				conn.Send <- message
			}

		}
	}

}

// WsHandler TestHandler socket 连接 中间件 作用:升级协议,用户验证,自定义信息等
func WsHandler(c *gin.Context) {
	conn, err := (&websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true
		}}).Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		http.NotFound(c.Writer, c.Request)
		return
	}
	//可以添加用户信息验证
	client := &Client{
		ID:     createUUID(),
		Socket: conn,
		Send:   make(chan []byte),
	}
	Manager.Register <- client
	go client.Read()
	go client.Write()
}

func SaveConfToLocal() {
	// 将Person对象编码为JSON格式
	jsonData, err := json.Marshal(AppConfigData)
	if err != nil {
		log.Fatal(err)
	}

	// 将编码后的JSON数据写入到文件中
	err = ioutil.WriteFile("config.json", jsonData, 0644)
	if err != nil {
		log.Fatal(err)
	}
}

func GetLoaclConfig() {
	// 读取JSON文件
	file, err := ioutil.ReadFile("config.json")
	if err != nil {
		log.Fatal(err)
	}

	// 将JSON文件中的数据作为Person对象解码
	err = json.Unmarshal([]byte(file), &AppConfigData)
	if err != nil {
		log.Fatal(err)
	}
}

config.json

{"data":"1233456","cron":"*/3 * * * * *","url":"wss://localhost:8080/echo","enable":false}

main.go

package main

import (
	_ "WebSocketTool/docs" //依赖项必须导入
	"crypto/tls"
	"encoding/json"
	"github.com/gin-gonic/gin"
	swaggerFiles "github.com/swaggo/files"
	"github.com/swaggo/gin-swagger"
	"io/ioutil"
	"log"
	"net/http"
)

// @Summary 开启定时发送消息
// @title Swagger Example API
// @version 0.0.1
// @description  This is a sample server Petstore server.
// @BasePath /
// @Host 127.0.0.1:8081
// @Accept json
// @Produce  json
// @Param params body	AppConfig true "开启定时发送配置"
// @Success 200 {string} json "{"code":200,"data":"name","msg":"ok"}"
// @Router /enableTiming [post]
func enableTiming(c *gin.Context) {
	bodyByts, err := ioutil.ReadAll(c.Request.Body)
	if err != nil {
		// 返回错误信息
		c.String(http.StatusBadRequest, err.Error())
		// 执行退出
		c.Abort()
	}
	log.Printf("获取到的请求参数：%v", string(bodyByts))

	// 解码 JSON 内容
	err = json.Unmarshal(bodyByts, &AppConfigData)
	if err != nil {
		panic(err)
	}
	log.Printf("格式化后的请求参数：%v", AppConfigData)
	// 返回的 code 和 对应的参数星系
	SaveConfToLocal()
	c.JSON(200, gin.H{
		"code": http.StatusOK,
		"data": AppConfigData,
		"msg":  "success",
	})
}

// @Summary 打印测试功能
// @title Swagger Example API
// @version 0.0.1
// @description  This is a sample server Petstore server.
// @BasePath /
// @Host 127.0.0.1:8081
// @Produce  json
// @Param name query string true "Name"
// @Success 200 {string} json "{"code":200,"data":"name","msg":"ok"}"
// @Router /print [get]
func Print(c *gin.Context) {
	var (
		name string
	)
	name = c.Query("name")
	c.JSON(http.StatusOK, gin.H{
		"code": http.StatusOK,
		"msg":  "success",
		"data": name,
	})
}

// @Summary 获取默认配置
// @title Swagger Example API
// @version 0.0.1
// @description  This is a sample server Petstore server.
// @BasePath /
// @Produce  json
// @Success 200 {string} json "{"code":200,"data":"name","msg":"ok"}"
// @Router /getConfig [get]
func getConfig(c *gin.Context) {
	GetLoaclConfig()
	c.JSON(http.StatusOK, gin.H{
		"code": http.StatusOK,
		"msg":  "success",
		"data": AppConfigData,
	})
}

func main() {
	gin.SetMode(gin.ReleaseMode) //线上环境
	go Manager.Start()
	router := gin.Default()

	cert := "fullchain.pem"
	key := "privkey.key"
	//
	config := &tls.Config{
		MinVersion:   tls.VersionTLS12,
		Certificates: []tls.Certificate{},
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(cert, key)
	if err != nil {
		log.Fatal(err)
	}

	// 解析HTML模板
	// 加载静态文件
	router.Static("/js", "templates/js")
	router.Static("/css", "templates/css")
	router.LoadHTMLFiles("templates/index.html")
	// 定义路由和处理函数
	router.GET("/", func(c *gin.Context) {
		c.HTML(http.StatusOK, "index.html", gin.H{
			"title": "Home",
		})
	})
	router.GET("/print", Print)
	router.GET("/getConfig", getConfig)
	router.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))
	router.GET("/echo", WsHandler)
	router.GET("/pong", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"message": "pong",
		})
	})
	router.POST("/enableTiming", enableTiming)

	// 使用TLS配置初始化HTTP路由
	server := &http.Server{
		Addr:      ":8080", //listen and serve on 0.0.0.0:8080
		Handler:   router,
		TLSConfig: config,
	}

	//启动HTTP服务器
	if err := server.ListenAndServeTLS("", ""); err != nil {
		log.Fatal(err)
	}
	//if err := server.ListenAndServe(); err != nil {
	//	log.Fatal(err)
	//}

}

v1版本index.html

<!DOCTYPE html>
<html>
<head>
    <title>WebSocket广播</title>
    <meta charset="utf-8">
    <!-- 引入 layui 样式文件 -->
    <link rel="stylesheet" href="https://www.layuicdn.com/layui/css/layui.css">
    <!-- 引入 jQuery -->
    <script src="https://code.jquery.com/jquery-3.5.1.min.js"></script>
    <!-- 引入 layui JavaScript 库 -->
    <script src="https://www.layuicdn.com/layui/layui.js"></script>
</head>
<body>
<div class="layui-container">
    <form class="layui-form" action="" lay-filter="myform">
        <div class="layui-form-item">
            <label class="layui-form-label">报文：</label>
            <div class="layui-input-block">
                <textarea id="data" placeholder="请输入内容" class="layui-textarea" name="data"></textarea>
            </div>
        </div>
        <div class="layui-form-item">
            <label class="layui-form-label">间隔：</label>
            <div class="layui-input-block">
                <input type="text" name="interval" placeholder="请输入间隔时间" value="{{.TimeInterval}}"
                       autocomplete="off" class="layui-input">
            </div>
        </div>
        <div class="layui-form-item">
            <div class="layui-input-block">
                <button class="layui-btn layui-btn-disabled" lay-submit lay-filter="form-submit">提交</button>
                <button type="reset" class="layui-btn layui-btn-primary">重置</button>
            </div>
        </div>
        <div class="layui-form-item">
            <label class="layui-form-label">WSS：</label>
            <div class="layui-input-block">
                <input type="text" name="ws" lay-verify="title" autocomplete="off" placeholder="请输入WS地址"
                       class="layui-input" value="{{.Url}}">
            </div>
        </div>
        <div class="layui-form-item">
            <div class="layui-input-block">
                <button class="layui-btn" lay-submit lay-filter="form-connect">连接</button>
                <button class="layui-btn layui-btn-normal" lay-submit lay-filter="form-send">发送</button>
                <button class="layui-btn layui-btn-primary" lay-submit lay-filter="form-disconnect">断开连接</button>
                <button id="clean_log" type="button" class="layui-btn layui-btn-primary">清空日志</button>
            </div>
        </div>
        <div class="layui-form-item layui-form-text">
            <label class="layui-form-label">接受报文：</label>
            <div class="layui-input-block">
                <textarea style="height: 500px;color: forestgreen;" id="response" name="recvmessage" placeholder="请输入内容"
                          class="layui-textarea"></textarea>
            </div>
        </div>

    </form>
</div>

<script>
    layui.use(['form', 'jquery', 'layer'], function () {
        const form = layui.form;
        var $ = layui.$;
        var layer = layui.layer;
        let ws;

        // 赋值
        $('#data').val("{{.Data}}");


        // 监听表单提交事件
        form.on('submit(form-submit)', function (data) {
            console.log('submit form', data.field);
            // 这里可以将表单数据提交给后端
            $.ajax({
                url: "/save",    //请求的url地址
                dataType: "json",   //返回格式为json
                async: true,//请求是否异步，默认为异步，这也是ajax重要特性
                data: data.field,    //参数值
                type: "GET",   //请求方式
                beforeSend: function () {
                    //请求前的处理
                },
                success: function (req) {
                    //请求成功时处理
                },
                complete: function () {
                    //请求完成的处理
                },
                error: function () {
                    //请求出错处理
                }
            });
            return false;
        });

        // 连接 WebSocket
        form.on('submit(form-connect)', function (data) {
            const url = data.field.ws;
            ws = new WebSocket(url);
            ws.onopen = function () {
                log('连接成功：' + url);
            };
            ws.onmessage = function (event) {
                const message = event.data;
                console.log('收到消息：' + message);
                console.log('文本域的值：' + form.val("myform").recvmessage);
                var val = form.val("myform").recvmessage + "\n" + new Date().toLocaleString() + "\n" + message
                form.val("myform", {recvmessage: val})
            };
            ws.onclose = function () {
                log('连接已关闭：' + url);
            };
            return false;
        });

        // 发送消息
        form.on('submit(form-send)', function (data) {
            const message = data.field.data;
            if (ws && ws.readyState === WebSocket.OPEN) {
                ws.send(message);
                console.log('发送消息：' + message);
            }
            return false;
        });

        // 断开连接
        form.on('submit(form-disconnect)', function () {
            if (ws) {
                ws.close();
                console.log('已断开连接');
            }
            return false;
        });

        $("#clean_log").click(function () {
            $('#response').val('');
        });

        // 清空日志


        // 打印日志
        function log(message) {
            $('textarea[name=recvmessage]').val(function (index, text) {
                return text + '\n' + message;
            });
        }

    });
</script>
</body>
</html>

v2版本Vue代码：

<template>
  <el-row>
    <el-col :span="6"><div class="grid-content bg-purple"></div></el-col>
    <el-col :span="12">
      <el-form ref="form" :model="form" label-width="100px">
        <el-form-item label="订阅地址">
          <el-input v-model="form.url"></el-input>
        </el-form-item>
        <el-form-item label="报文">
          <el-input type="textarea" v-model="form.data" rows="8"></el-input>
        </el-form-item>
        <el-form-item label="操作">
          <el-row>
            <el-button
              type="success"
              style="width: 88px"
              @click="handleConnectWs"
              >连接</el-button
            >
            <el-button type="primary" style="width: 88px" @click="handleSendMsg"
              >发送</el-button
            >
            <el-button type="danger" style="width: 88px" @click="handleCloseWs">
              断开</el-button
            >
            <el-button type="info" style="width: 88px" @click="handleClearLog"
              >清空日志</el-button
            >
          </el-row>
        </el-form-item>
        <el-form-item label="开启定时发送">
          <el-input
            type="text"
            placeholder="请输入内容"
            v-model="form.cron"
            style="width: 200px"
          >
          </el-input>
          <!-- <el-input-number
            v-model="form.interval"
            @change="handleChange"
            :min="1"
            :max="10"
            label="描述文字"
          ></el-input-number> -->
          <el-switch
            v-model="form.enable"
            active-color="#13ce66"
            style="margin-left: 20px; margin-right: 20px"
          >
          </el-switch>
          <el-button type="primary" round @click="handleUpdate"
            >更新配置</el-button
          >
        </el-form-item>
        <el-form-item label="接受报文">
          <el-input
            id="scroll_text"
            ref="log"
            type="textarea"
            v-model="recvData"
            rows="20"
          ></el-input>
        </el-form-item>
      </el-form>
    </el-col>
    <el-col :span="6">
      <el-result
        :icon="connState.msg"
        :title="connState.title"
        :subTitle="connState.subTitle"
      >
      </el-result>
    </el-col>
  </el-row>
</template>

<script>
export default {
  name: "ConfigPanel",
  data() {
    return {
      connState: {
        msg: "warning",
        title: "未连接！",
        subTitle: "请先建立连接",
      },
      form: {
        url: "ws://127.0.0.1:8081/echo",
        data: "",
        enable: false,
        cron: "*/3 * * * * *",
      },
      websocket: null,
      recvData: "",
    };
  },
  mounted() {
    this.initData();
  },
  watch: {
    recvData(newValue) {
      this.$nextTick(() => {
        const textarea = this.$refs.log;
        if (newValue !== textarea.value) {
          textarea.value = newValue; // 更新textarea的值
        }
        textarea.scrollTop = textarea.scrollHeight;
      });
    },
  },
  methods: {
    initData() {
      this.$api.getConfig().then((res) => {
        console.log(res.data);
        const { url, data, enable, cron } = { ...res.data };
        Object.assign(this.form, { url, data, enable, cron });
        console.log(this.form);
      });
    },

    /**
     * 连接wS点击事件
     */
    handleConnectWs() {
      this.createWebsocket();
    },

    /**
     * 发送消息点击事件
     */
    handleSendMsg() {
      console.log(this.websocket);
      if (this.websocket == null) {
        this.$notify({
          type: "warning",
          title: "WbSocket未连接！",
          message: "WbSocket未连接，请先进行连接！",
        });
        return;
      }
      this.websocket.send(this.form.data);
    },
    /**
     * 关闭WS连接点击事件
     */
    handleCloseWs() {
      if (this.websocket == null) {
        this.$notify({
          type: "warning",
          title: "WbSocket未连接！",
          message: "WbSocket未连接，请先进行连接！",
        });
        return;
      }
      this.websocket.close();
    },
    /**
     * 清空日志点击事件
     */
    handleClearLog() {
      this.recvData = "";
    },
    /**
     * 更新配置
     */
    handleUpdate() {
      console.log(this.form);
      this.$api.enableTiming({ ...this.form }).then((res) => {
        console.log(res);
        if (res.code === 200) {
          this.$notify({
            title: "成功",
            message: "更新配置成功，重启生效！",
            type: "success",
          });
        }
      });
    },
    createWebsocket() {
      if (this.websocket != null) {
        return;
      }
      this.websocket = new WebSocket(this.form.url);
      // 连接发生错误的回调方法
      this.websocket.onerror = this.websocketOnerror;
      // 连接成功建立的回调方法
      this.websocket.onopen = this.websocketOnopen;
      // 接收到消息的回调方法
      this.websocket.onmessage = this.websocketOnmessage;
      // 连接关闭的回调方法
      this.websocket.onclose = this.websocketOnclose;
      // 监听窗口关闭事件，当窗口关闭时，主动去关闭websocket连接，防止连接还没断开就关闭窗口，server端会抛异常。
      this.websocket.onbeforeunload = this.websocketOnbeforeunload;
    },
    // 连接发生错误的回调方法
    websocketOnerror() {
      console.log("连接发生错误的回调方法");
      this.$notify.error({
        title: "错误",
        message: "连接WS" + this.form.url + "失败",
      });
      this.websocket = null;
      this.connState = {
        msg: "error",
        title: "连接失败",
        subTitle: "请先建立连接",
      };
    },
    // 连接成功建立的回调方法
    websocketOnopen() {
      console.log("连接成功建立的回调方法");
      this.$notify({
        title: "成功",
        message: "建立连接成功！",
        type: "success",
      });
      this.connState = {
        msg: "success",
        title: "连接成功",
        subTitle: "",
      };
    },
    // 接收到消息的回调方法
    websocketOnmessage(event) {
      console.log("接收到消息的回调方法");
      const message = event.data;
      this.recvData +=
        this.formateDate(new Date(), "yyyy-MM-dd hh:mm:ss") +
        "\n" +
        message +
        "\n";
      this.handleInputChange;
    },
    // 连接关闭的回调方法
    websocketOnclose() {
      this.$notify({
        title: "WbSocket已关闭！",
        message: "WbSocket已关闭!",
      });
      console.log("连接关闭的回调方法");
      this.recvData +=
        this.formateDate(new Date(), "yyyy-MM-dd hh:mm:ss") +
        "\n" +
        "WS连接关闭" +
        "\n";
      this.connState = {
        msg: "warning",
        title: "未连接！",
        subTitle: "请先建立连接",
      };
      this.websocket = null;
    },
    // 监听窗口关闭事件，当窗口关闭时，主动去关闭websocket连接，防止连接还没断开就关闭窗口，server端会抛异常
    websocketOnbeforeunload() {
      this.closeWebSocket();
      console.log(
        "监听窗口关闭事件，当窗口关闭时，主动去关闭websocket连接，防止连接还没断开就关闭窗口，server端会抛异常"
      );
    },
    // 关闭WebSocket连接
    closeWebSocket() {
      this.websocket.close();
    },
    onSubmit() {
      console.log("submit!");
    },

    formateDate(date, fmt) {
      if (/(y+)/.test(fmt)) {
        fmt = fmt.replace(
          RegExp.$1,
          (date.getFullYear() + "").substr(4 - RegExp.$1.length)
        );
      }
      let o = {
        "M+": date.getMonth() + 1,
        "d+": date.getDate(),
        "h+": date.getHours(),
        "m+": date.getMinutes(),
        "s+": date.getSeconds(),
      };
      for (let k in o) {
        if (new RegExp(`(${k})`).test(fmt)) {
          let str = o[k] + "";
          fmt = fmt.replace(
            RegExp.$1,
            RegExp.$1.length === 1 ? str : this.padLeftZero(str)
          );
        }
      }
      return fmt;
    },
    // 左边补0函数
    padLeftZero(str) {
      return ("00" + str).substr(str.length);
    },
  },
};
</script>

打包：

go build -o D:\Go_WorkSpace\test2\bin\WebSocketTool.exe test2/src/main #gosetup