golang的socket的安全(连接安全)

对于一个socket服务,没有任何安全措施肯定是不对的，这样就会造成任何人都可以来连接，进而造成不可预知的风险。现在我们来实现一个有安全认证的一个服务.

package main

/*
* 安全认证1,只实现了认证，仍存在风险
 */

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"
)

//保存连接
var connList sync.Map

//消息结构{"from":"aaa","to":"bbb","body":"aaa","cmd":"sendMsg"}
type UserMsg struct {
	From string `json:"from"`
	To   string `json:"to"`
	Body string `json:"body"`
	Cmd  string `json:"cmd"`
}

//随机字符串
func randName() string {
	var letterRunes = []rune("123456789")
	rand.Seed(time.Now().UnixNano())
	b := make([]rune, 6)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func DealConn(conn net.Conn) {
	// 处理完关闭连接
	defer conn.Close()
	fmt.Println("new conn\n")
	conn.Write([]byte("hello\n"))
	//安全认证
	reader := bufio.NewReader(conn)
	str, err := reader.ReadString('\n')
	if err == nil {
		str = strings.TrimSpace(str)
		var strStruct UserMsg
		err := json.Unmarshal([]byte(str), &strStruct)
		if err != nil {
			fmt.Printf("msg json error, err:%v\n", err)
			return
		} else {
			if strStruct.Cmd == "auth" {
				if strStruct.Body != "123456" {
					fmt.Println("not auth")
					return
				}
			} else {
				fmt.Println("not auth")
				return
			}
		}
	} else {
		return
	}

	name := randName()
	connList.Store(name, conn)
	conn.Write([]byte("you is  : " + name + "\r\n"))

	// 针对当前连接做发送和接受操作
	for {
		reader := bufio.NewReader(conn)
		// 读取字符串, 直到碰到回车返回
		str, err := reader.ReadString('\n')
		// 数据读取正确
		if err == nil {
			// 去掉字符串尾部的回车
			str = strings.TrimSpace(str)
			var strStruct UserMsg
			err := json.Unmarshal([]byte(str), &strStruct)
			if err != nil {
				fmt.Printf("msg json error, err:%v\n", err)
			} else {
				to := strStruct.To
				cmd := strStruct.Cmd
				if cmd == "sendMsg" {
					toUser, ok := connList.Load(to)
					if !ok {
						fmt.Println("to user is not\n")
					} else {
						toUserConn, ok := toUser.(net.Conn)
						if !ok {
							fmt.Println("to user type wrong\n")
						} else {
							toUserConn.Write([]byte(strStruct.From + " say : " + strStruct.Body + "\r\n"))
						}
					}
				} else {
					fmt.Println("unkonw msg")
				}
			}
		} else {
			//删除保存的连接
			connList.Delete(name)
			fmt.Println("conn close\n")
			break
		}
	}
}

func main() {
	// 建立tcp 服务
	listen, err := net.Listen("tcp", "0.0.0.0:1215")
	if err != nil {
		fmt.Printf("listen failed, err:%v\n", err)
		return
	}

	for {
		// 等待客户端建立连接
		fmt.Printf("accept conn ...... \n")
		conn, err := listen.Accept()
		if err != nil {
			fmt.Printf("accept failed, err:%v\n", err)
			continue
		}

		// 启动一个单独的 goroutine 去处理连接
		go DealConn(conn)
	}
}

本代码基本实现了一个连接认证的逻辑，没有认证的连接是无法进行操作的，但是仍然存在一定的安全隐患，如果一个连接来了之后不进行认证，一直呆在那里，如果这样的空闲连接特别的多，也会造成服务资源的消耗和浪费，接下来我们看看怎么进行优化

package main

/*
* 安全认证2,实现了认证加超时
 */

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"strings"
	"sync"
	"time"
)

//保存连接
var connList sync.Map

//消息结构{"from":"aaa","to":"bbb","body":"aaa","cmd":"sendMsg"}
type UserMsg struct {
	From string `json:"from"`
	To   string `json:"to"`
	Body string `json:"body"`
	Cmd  string `json:"cmd"`
}

//随机字符串
func randName() string {
	var letterRunes = []rune("123456789")
	rand.Seed(time.Now().UnixNano())
	b := make([]rune, 6)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}

func DealConn(conn net.Conn) {
	// 处理完关闭连接
	defer conn.Close()
	fmt.Println("new conn\n")
	conn.Write([]byte("hello\n"))
	//安全认证+读取超时
	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	reader := bufio.NewReader(conn)
	str, err := reader.ReadString('\n')
	if err == nil {
		str = strings.TrimSpace(str)
		var strStruct UserMsg
		err := json.Unmarshal([]byte(str), &strStruct)
		if err != nil {
			fmt.Printf("msg json error, err:%v\n", err)
			return
		} else {
			if strStruct.Cmd == "auth" {
				if strStruct.Body != "123456" {
					fmt.Println("not auth")
					return
				}
			} else {
				fmt.Println("not auth")
				return
			}
		}
	} else {
		return
	}

	name := randName()
	connList.Store(name, conn)
	conn.Write([]byte("you is  : " + name + "\r\n"))

	// 针对当前连接做发送和接受操作
	for {
		conn.SetReadDeadline(time.Now().Add(3600 * time.Second))
		reader := bufio.NewReader(conn)
		// 读取字符串, 直到碰到回车返回
		str, err := reader.ReadString('\n')
		// 数据读取正确
		if err == nil {
			// 去掉字符串尾部的回车
			str = strings.TrimSpace(str)
			var strStruct UserMsg
			err := json.Unmarshal([]byte(str), &strStruct)
			if err != nil {
				fmt.Printf("msg json error, err:%v\n", err)
			} else {
				to := strStruct.To
				cmd := strStruct.Cmd
				if cmd == "sendMsg" {
					toUser, ok := connList.Load(to)
					if !ok {
						fmt.Println("to user is not\n")
					} else {
						toUserConn, ok := toUser.(net.Conn)
						if !ok {
							fmt.Println("to user type wrong\n")
						} else {
							toUserConn.Write([]byte(strStruct.From + " say : " + strStruct.Body + "\r\n"))
						}
					}
				} else {
					fmt.Println("unkonw msg")
				}
			}
		} else {
			//删除保存的连接
			connList.Delete(name)
			fmt.Println("conn close\n")
			break
		}
	}
}

func main() {
	// 建立tcp 服务
	listen, err := net.Listen("tcp", "0.0.0.0:1215")
	if err != nil {
		fmt.Printf("listen failed, err:%v\n", err)
		return
	}

	for {
		// 等待客户端建立连接
		fmt.Printf("accept conn ...... \n")
		conn, err := listen.Accept()
		if err != nil {
			fmt.Printf("accept failed, err:%v\n", err)
			continue
		}

		// 启动一个单独的 goroutine 去处理连接
		go DealConn(conn)
	}
}

以上主要的逻辑就是在认证读取的时候增加了超时进而主动关闭连接，防止空闲连接太多.