概括

今天了解了一下websocket的简单使用，并用它别写了一个小的聊天室程序，话不多说，直接上代码，记录一下学习成果

服务端

var UP = websocket.Upgrader{
   ReadBufferSize: 1024,
   WriteBufferSize: 1024,
}

//将用户名与连接以k-v进行存储
var conns = make(map[string]*websocket.Conn)
//存储账号密码
var user = make(map[string]string)

func handle(w http.ResponseWriter,r *http.Request)  {
   var (
      flag bool
      username string
   )

   conn, err := UP.Upgrade(w, r, nil)
   defer conn.Close()
   username = r.RemoteAddr
   if err != nil {
      fmt.Println(err.Error())
      return
   }

   log.Println(r.RemoteAddr,"接入连接")

   for !flag {
      conn.WriteMessage(websocket.TextMessage,[]byte("---------欢迎来到聊天室----------"))
      conn.WriteMessage(websocket.TextMessage,[]byte("-------请输入数字选择功能--------"))
      conn.WriteMessage(websocket.TextMessage,[]byte("-----1.注册  2.登录  3.退出-----"))
      _, bytes, err := conn.ReadMessage()
      if err != nil {
         return
      }
      choice := string(bytes)

      switch choice {
         case "1":
            username,flag = signin(conn)
            break
         case "2":
            username,flag = login(conn)
            break
         case "3":
            log.Println(username,"断开了连接")
            return
      default:
         conn.WriteMessage(websocket.TextMessage,[]byte("输入错误请重试"))
         flag = false
      }
   }

   for  {
      _, bytes, err := conn.ReadMessage()
      if err != nil {
         break
      }
      for _,conn := range conns{
         conn.WriteMessage(websocket.TextMessage,[]byte(username + "：" + string(bytes)))
      }
      fmt.Println(string(bytes))
   }
   for _,conn := range conns{
      conn.WriteMessage(websocket.TextMessage,[]byte(username + "退出了聊天室"))
   }
   log.Println(r.RemoteAddr,"断开了连接")
}

//注册
func signin(conn *websocket.Conn) (username string,flag bool) {
   username, password := input(conn)
   user[username] = password
   conns[username] = conn
   conn.WriteMessage(websocket.TextMessage,[]byte("注册成功..."))
   for _,conn:= range conns{
      conn.WriteMessage(websocket.TextMessage,[]byte(username + "成功进入聊天室...."))
   }
   return username,true
}

//登录
func login(conn *websocket.Conn) (username string,flag bool){
   username, password := input(conn)
   flag1 := cheak(username, password)
   flag2 := conns[username] == nil
   if flag1 && flag2 {
      conns[username] = conn
      flag = true
      for _,conn := range conns{
         conn.WriteMessage(websocket.TextMessage,[]byte(username + "成功进入聊天室...."))
      }
   }else if flag1 {
      conn.WriteMessage(websocket.TextMessage,[]byte("该用户已登录"))
      flag = false
      return
   }else {
      conn.WriteMessage(websocket.TextMessage,[]byte("用户名或密码错误"))
      flag = false
      return
   }

   return
}

//得到用户账号密码
func input(conn *websocket.Conn) (username ,password string){
   conn.WriteMessage(websocket.TextMessage,[]byte("请输入账号和密码"))
   conn.WriteMessage(websocket.TextMessage,[]byte("账号："))
   for i := 0;i < 2;i++ { //接收账号与密码，所以接收两次就行
      _, bytes, err := conn.ReadMessage()
      if err != nil {
         break
      }
      if i == 0 { //第一次接收账号，并给出输入密码的提示
         username = string(bytes)
         conn.WriteMessage(websocket.TextMessage, []byte("密码："))
      } else {
         password = string(bytes) //接收密码
      }
   }
   return
}

//校验用户密码
func cheak(username string, password string) bool { //与user map中进行校验，用户是否存在
   if user[username] == password {
      return true
   }else {
      return false
   }
}

func main() {
   http.HandleFunc("/",handle)
   http.ListenAndServe(":8080",nil)
}

服务端是用了 github.com/gorilla/websocket 这个包，写的不是十分完善，比如用户密码存储不是持久的，当服务端关闭时，数据全部丢失，下次再改进哈哈

客户端

func main() {
   dialer := websocket.Dialer{}
   dial, _, err := dialer.Dial("ws://127.0.0.1:8080", nil)
   defer dial.Close()
   if err != nil {
      return
   }
   go send(dial)
   for  {
      _, bytes, err := dial.ReadMessage()
      if err != nil {
         break
      }

      fmt.Println(string(bytes))
   }
}

func send(dial *websocket.Conn) {
   for  {
      reader := bufio.NewReader(os.Stdin)
      line, _, err := reader.ReadLine()
      if err != nil {
         return
      }
      dial.WriteMessage(websocket.TextMessage,line)
   }

}

本机测试部分效果如下：

本文正在参加技术专题18期-聊聊Go语言框架