Golang TCPConn.SetReadDeadline方法代码示例

//
// A TCP 'echo' example.
// Demonstrates using a timeout, and a 'graceful' shutdown if one occurs.
// 'Tested' using 'telnet localhost 45678'
//
func runReads(tcpConn *net.TCPConn) bool {
	br := bufio.NewReader(tcpConn)
	for {
		// Set a timeout value, which needs to be set before each and every read.
		d := time.Duration(30 * 1e9) // 30 seconds
		w := time.Now()              // from now
		w = w.Add(d)
		tcpConn.SetReadDeadline(w) // Set the deadline
		//
		buffer, err := br.ReadBytes('\n') // '\n' is delimiter
		// If the read times out, this prints something like:
		// Error = read tcp 127.0.0.1:57609: resource temporarily unavailable
		if err != nil {
			fmt.Printf("Error = %v\n", err)
			return false
			// panic("wtf04")
		}
		//
		fmt.Printf("Bytes Read: %d\n", len(buffer))
		var data = string(buffer)
		fmt.Printf("Data Read: |%q|\n", data)

		// This is now an 'echo' example.
		out := "echo: " + data
		tcpConn.Write([]byte(out))

		// The \r in this data from telnet is a bit surprising ...
		if data == "quit\r\n" {
			fmt.Println("Breaking....")
			break
		}
	}
	return true
}

func (srv *Server) readTCP(conn *net.TCPConn, timeout time.Duration) ([]byte, error) {
	conn.SetReadDeadline(time.Now().Add(timeout))
	l := make([]byte, 2)
	n, err := conn.Read(l)
	if err != nil || n != 2 {
		if err != nil {
			return nil, err
		}
		return nil, ErrShortRead
	}
	length, _ := unpackUint16(l, 0)
	if length == 0 {
		return nil, ErrShortRead
	}
	m := make([]byte, int(length))
	n, err = conn.Read(m[:int(length)])
	if err != nil || n == 0 {
		if err != nil {
			return nil, err
		}
		return nil, ErrShortRead
	}
	i := n
	for i < int(length) {
		j, err := conn.Read(m[i:int(length)])
		if err != nil {
			return nil, err
		}
		i += j
	}
	n = i
	m = m[:n]
	return m, nil
}

func (server *TcpServer) Start() (err error) {
	for {
		var conn *net.TCPConn
		if conn, err = server.TCPListener.AcceptTCP(); err != nil {
			return err
		}
		if server.keepAlive != nil {
			if err := conn.SetKeepAlive(server.keepAlive.(bool)); err != nil {
				return err
			}
		}
		if server.keepAlivePeriod != nil {
			if kap, ok := (net.Conn(conn)).(iKeepAlivePeriod); ok {
				if err := kap.SetKeepAlivePeriod(server.keepAlivePeriod.(time.Duration)); err != nil {
					return err
				}
			}
		}
		if server.linger != nil {
			if err := conn.SetLinger(server.linger.(int)); err != nil {
				return err
			}
		}
		if server.noDelay != nil {
			if err := conn.SetNoDelay(server.noDelay.(bool)); err != nil {
				return err
			}
		}
		if server.readBuffer != nil {
			if err := conn.SetReadBuffer(server.readBuffer.(int)); err != nil {
				return err
			}
		}
		if server.writerBuffer != nil {
			if err := conn.SetWriteBuffer(server.writerBuffer.(int)); err != nil {
				return err
			}
		}
		if server.deadline != nil {
			if err := conn.SetDeadline(server.deadline.(time.Time)); err != nil {
				return err
			}
		}
		if server.readDeadline != nil {
			if err := conn.SetReadDeadline(server.readDeadline.(time.Time)); err != nil {
				return err
			}
		}
		if server.writerDeadline != nil {
			if err := conn.SetWriteDeadline(server.writerDeadline.(time.Time)); err != nil {
				return err
			}
		}
		if server.config != nil {
			server.ServeTCP(tls.Client(conn, server.config))
		} else {
			server.ServeTCP(conn)
		}
	}
}

func (t *Duo) Pipe(conn *net.TCPConn) {
	fd := t.Fd(conn)
	defer func() {
		t.Logger.Printf("disconnected: %s\n", fd)
		conn.Close()
		delete(t.Conn, fd)
	}()

	// Save in map
	t.Conn[fd] = conn

	// Read data
	reader := bufio.NewReader(conn)
	for {
		body, err := t.Unpack(reader)
		if err != nil {
			if err == io.EOF {
				return
			}
			if nerr, ok := err.(net.Error); ok && nerr.Timeout() {
				return
			}
			t.Logger.Print(err)
		}

		conn.SetReadDeadline(time.Now().Add(20 * time.Second))

		if len(body) == 0 {
			continue
		}

		t.handler(conn, body)
		reader.Reset(conn)
	}
}

// 协议封装读取
func TCPConnRead(conn *net.TCPConn) ([]byte, error) {
	conn.SetReadDeadline(time.Now().Add(time.Second * 30))
	result := bytes.NewBuffer(nil)
	data := make([]byte, 4)
	num, err := conn.Read(data)
	if err != nil || num != 4 {
		if err == nil {
			err = errors.New("length read error")
		}
		return nil, err
	}
	result.Write(data[0:num])
	var length int32
	err = binary.Read(result, binary.LittleEndian, &length)
	if err != nil {
		return nil, err
	}
	if length > MAX_PACKET_SIZE {
		return nil, errors.New("too large packet! packet size should less than 1M")
	}
	data = make([]byte, length)
	result = bytes.NewBuffer(nil)
	num, err = io.ReadFull(conn, data)
	if err != nil {
		return nil, err
	}
	result.Write(data[0:num])
	return result.Bytes(), nil
}

func receiveTcpConn(conn *net.TCPConn, rc TcpReceiveCaller) {
	for {
		buf := bufio.NewReader(conn)
		line, err := buf.ReadBytes('\n')
		if err != nil {
			// remove client
			addr := conn.RemoteAddr().String()
			if v, ok := clients[addr]; ok {
				uid := v.UserId
				delete(clients, addr)
				addr2 := users[uid]
				if strings.Index(addr2, "$") == -1 {
					delete(users, uid)
				} else {
					users[uid] = strings.Replace(strings.Replace(addr2, addr, "", 1), "$$", "$", -1)
				}
			}
			printf(true, "[ CLIENT][ DISCONN] - IP : %s disconnect!active clients : %d",
				conn.RemoteAddr().String(), len(clients))
			break

		}
		if d, err := rc(conn, line[:len(line)-1]); err != nil { // remove '\n'
			conn.Write([]byte("error$" + err.Error()))
		} else if d != nil {
			conn.Write(d)
		}
		conn.Write([]byte("\n"))
		conn.SetReadDeadline(time.Now().Add(ReadDeadLine)) // discount after 5m
	}
}

func session(conn *net.TCPConn) {
	fmt.Println("here")
	var buf [2048]byte
	code := 0
	for {
		t := time.Now().Add(time.Millisecond * 100)
		conn.SetReadDeadline(t)
		n, err := conn.Read(buf[:])
		e, ok := err.(net.Error)

		if err != nil && ok && !e.Timeout() {
			fmt.Println(err)
			break
		}

		if n > 0 {
			process(conn, buf[:n])
		} else {
			msg := fmt.Sprintf("%v", code)
			code++
			conn.Write([]byte(msg))
		}
	}
	fmt.Println("session ended")
}

func (dc *DefaultConnector) receive(conn *net.TCPConn) ([]byte, error) {
	err := conn.SetReadDeadline(time.Now().Add(dc.config.ReadTimeout))
	if err != nil {
		return nil, err
	}

	header := make([]byte, 8)
	_, err = io.ReadFull(conn, header)
	if err != nil {
		return nil, err
	}

	decoder := NewBinaryDecoder(header)
	length, err := decoder.GetInt32()
	if err != nil {
		return nil, err
	}
	response := make([]byte, length-4)
	_, err = io.ReadFull(conn, response)
	if err != nil {
		return nil, err
	}

	return response, nil
}

// read a reply to a buffer based on the expected message type
// return error if reply message has different type of command than expected
func readHelloReply(conn *net.TCPConn) (proto.Message, error) {
	duration := time.Second * 10
	timeNow := time.Now()
	err := conn.SetReadDeadline(timeNow.Add(duration))
	if err != nil {
		TimeEncodedPrint("Cant set read timeout", err.Error())
		return nil, err
	}
	length := int32(0)
	err = binary.Read(conn, binary.LittleEndian, &length)
	if DEBUG_READING_MESSAGE {
		fmt.Println("TRYING TO READ MESSAGE LENGTH => ", length, " ERROR: ", err)
	}

	if err != nil {
		return nil, err
	}

	if BYPASS_CONNECTION_SERVER {
		tempBuf := make([]byte, 8)
		_, err = io.ReadFull(conn, tempBuf)
		if DEBUG_READING_MESSAGE {
			fmt.Println("Trying to read extra 8 bytes:", tempBuf, " PROBLEM: ", err)
		}

		if err != nil {
			return nil, err
		}
		length = length - 8
	}

	rbuf := make([]byte, length)
	n, err := io.ReadFull(conn, rbuf)

	cmd := int(rbuf[0])

	switch cmd {
	case S2C_HelloInfoResult_CMD:
		res := &Auth_S2C.HelloInfoResult{}
		err := proto.Unmarshal(rbuf[1:], res)
		if err != nil {
			log.Fatal(err)
		} else {
			return res, nil
		}

	case ERROR_ErrorInfo_CMD:
		res := &Auth_S2C.ErrorInfo{}
		err := proto.Unmarshal(rbuf[1:], res)
		if err != nil {
			log.Fatal(err)
		}
		msg := fmt.Sprint("Server returns error: ")
		return res, errors.New(msg)
	default:
		log.Fatal("Unexpected CMD: ", cmd, "length is: ", length, " number bytes read: ", n, " error: ", err)
	}
	return nil, nil

}

func pipeThenClose(src, dst *net.TCPConn, finishChannel chan bool) {
	defer func() {
		src.CloseRead()
		dst.CloseWrite()
		finishChannel <- true
	}()

	buf := asocks.GetBuffer()
	defer asocks.GiveBuffer(buf)

	for {
		src.SetReadDeadline(time.Now().Add(60 * time.Second))
		n, err := src.Read(buf)
		if n > 0 {
			data := buf[0:n]
			encodeData(data)
			if _, err := dst.Write(data); err != nil {
				break
			}
		}
		if err != nil {
			break
		}
	}
}

//处理连接请求
func (this *server) handlerConnection(conn *net.TCPConn) {

	defer logger.CatchException()

	logger.Infof("New connection coming ... IP=%s ", conn.RemoteAddr())

	conn.SetNoDelay(true)                                        //无延迟
	conn.SetKeepAlive(true)                                      //保持激活
	conn.SetReadBuffer(64 * 1024)                                //设置读缓冲区大小
	conn.SetWriteBuffer(64 * 1024)                               //设置写缓冲区大小
	conn.SetReadDeadline(time.Now().Add(30000000 * time.Second)) //设置读超时

	session := network.NewSession(conn)
	defer session.Close()

	for {
		msg, err := session.RecvMSG()
		if err != nil {
			logger.Infof("RecvMsgs IP=%s err=%v", conn.RemoteAddr(), err.Error())
			return
		}

		ret := this.ProcessMessage(session, msg)
		if ret == false {
			return
		}
	}
}

// Takes two existing TCP connections and copies all incoming traffic to the
// other channel. Returns only upon an error
// This function has some commented out code. This was me playing with TCP
// keepalive and the annoying missing API from go.
// It doesn't work so it's now kept simple and just disconnects after inactivity.
func (p *proxy) pipe(src, dst *net.TCPConn) {
	islocal := src == p.lconn

	// Detecting a broken/down connection in TCP is a little tricky as it turns
	// out:
	// http://tldp.org/HOWTO/TCP-Keepalive-HOWTO/usingkeepalive.html
	// $ tail /proc/sys/net/ipv4/tcp_keepalive_*
	// ==> /proc/sys/net/ipv4/tcp_keepalive_intvl <==
	// 75
	// ==> /proc/sys/net/ipv4/tcp_keepalive_probes <==
	// 9
	// ==> /proc/sys/net/ipv4/tcp_keepalive_time <==
	// 7200
	// Unfortunately GO sets both, interval and idle time :(
	// http://felixge.de/2014/08/26/tcp-keepalive-with-golang.html

	// src.SetKeepAlive(true)
	// src.SetKeepAlivePeriod(TCP_KEEP_ALIVE) // Bad don't set (read blog post above)
	// linuxEnableKeepAlive(src)

	// directional copy (16k buffer)
	buff := make([]byte, 0x3fff)
	// This loop is our main copy loop. It takes each packet and just shoves it
	// into the other side. This works for both ways
	for {
		// We will set a timeout so that we also drop clients after a while of
		// no activity.
		// Problem: This also dropps the connection even if sitting on an idle
		// shell. But who really leaves it hanging for days without any input/output?
		src.SetReadDeadline(time.Now().Add(TCP_READ_TIMEOUT))
		n, err := src.Read(buff)
		if err != nil {
			if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
				// Activate TCPs keepalive now which will close the connection
				// if the client is gone
				// TODO: Doesnt work:
				// linuxEnableKeepAlive(src)
				p.disconnect("Inactive client '%s'\n", err)
				return
			} else {
				p.disconnect("Read failed '%s'\n", err)
				return
			}
		}
		b := buff[:n]
		// write out result
		n, err = dst.Write(b)
		if err != nil {
			p.disconnect("Write failed '%s'\n", err)
			return
		}
		if islocal {
			p.sentBytes += uint64(n)
		} else {
			p.receivedBytes += uint64(n)
		}
	}
}

// read a reply to a buffer based on the expected message type
// return error if reply message has different type of command than expected
func readReply(conn *net.TCPConn) (proto.Message, error) {
	duration := time.Millisecond * 20
	timeNow := time.Now()
	err := conn.SetReadDeadline(timeNow.Add(duration))
	if err != nil {
		TimeEncodedPrint("Cant set read timeout", err.Error())
		return nil, err
	}
	length := int32(0)
	err = binary.Read(conn, binary.LittleEndian, &length)
	if err != nil {
		return nil, nil
	}

	if DEBUG_STEP {
		fmt.Println("RECEIVE MESSAGE LENGTH: ", length)
	}
	// now wait longer for the message data
	duration = time.Millisecond * 100
	timeNow = time.Now()
	err = conn.SetReadDeadline(timeNow.Add(duration))

	rbuf := make([]byte, length)
	io.ReadFull(conn, rbuf)
	if DEBUG_STEP {
		fmt.Println(rbuf)
	}

	cmd := int(rbuf[0])

	switch cmd {
	case S2C_HelloInfoResult_CMD:
		res := &Auth_S2C.HelloInfoResult{}
		err := proto.Unmarshal(rbuf[1:], res)
		if err != nil {
			log.Fatal(err)
		} else {
			return res, nil
		}

	case ERROR_ErrorInfo_CMD:
		res := &Auth_S2C.ErrorInfo{}
		err := proto.Unmarshal(rbuf[1:], res)
		if err != nil {
			log.Fatal(err)
		}
		msg := fmt.Sprint("Server returns error: ")
		return res, errors.New(msg)
	default:
		log.Fatal("Unexpected CMD: ", cmd)
	}
	return nil, nil

}

func (s Sender) sendData(data []byte, conn *net.TCPConn) bool {
	if len(data) == 0 {
		return true
	}

	if conn == nil {
		return false
	}
	/*
	   lenBuf := make([]byte, 4)
	   nData := len(data)
	   binary.PutUvarint(lenBuf, uint64(nData))
	   data = append(lenBuf, data...)
	*/

	st := time.Now()
	packId := tcp_pack.GetPackId(data)

	conn.SetDeadline(time.Now().Add(5 * time.Minute)) //设置超时
	loglib.Info(fmt.Sprintf("sender%d start sending pack:%s length:%d", s.id, packId, len(data)))
	n, err := conn.Write(data)
	ed := time.Now()
	loglib.Info(fmt.Sprintf("sender%d end sending pack:%s length:%d elapse:%s", s.id, packId, n, ed.Sub(st)))

	lib.CheckError(err)

	//写失败了就不用等应答了，肯定拿不到
	if err == nil {
		conn.SetReadDeadline(time.Now().Add(8 * time.Minute)) //设置超时
		time1 := time.Now()
		var temp []byte = make([]byte, 128)
		count, err := conn.Read(temp)
		if err == nil {
			loglib.Info(fmt.Sprintf("sender%d get anwser data len:%d for pack:%s elapse:%s", s.id, count, packId, time.Now().Sub(time1)))
		} else {
			loglib.Info(fmt.Sprintf("sender%d get anwser data len:%d for pack:%s elapse:%s, error:%s", s.id, count, packId, time.Now().Sub(time1), err.Error()))
		}

		temp = temp[:count]
		if string(temp) == "ok" { //发送成功
			return true
		} else if string(temp) == "wrong header" {
			//包头错误,丢弃
			loglib.Info(packId + " has wrong header, retry later!")
			return false
		} else { //发送失败
			//报警
			return false
		}
	} else {
		loglib.Warning(fmt.Sprintf("write pack %s error:%s", packId, err.Error()))
	}
	return false
}

/*
	keep routine for write and read
*/
func serverKeeper(conn *net.TCPConn) {
	// conn.SetKeepAlive(true)
	request := make([]byte, 1024) // set maxium request length to 128B to prevent flood attack
	defer func() {
		sendChan <- "&**& quit this server &**&" //退出协程并且删掉主对象
	}() // close connection before exit

	// send
	go func() {
		for {
			sendData := <-sendChan
			if sendData != "&**& quit this server &**&" {
				// fmt.Println("server sending...", sendData)
				_, err := conn.Write([]byte(sendData))
				checkError(err)
			} else {
				//崩溃退出关闭发送协程
				conn.Close()
				fmt.Println("quit server Keeper success")
				return
			}
		}
	}()

	//read
	for {
		conn.SetReadDeadline(time.Now().Add(2 * time.Minute)) // set 2 minutes timeout
		read_len, err := conn.Read(request)

		if err != nil {
			if err.Error() == "EOF" {
				fmt.Println("this is EOF")
				break
			}
			fmt.Println(err)
			fmt.Println(err.Error())
			//break
		}

		if read_len == 0 {
			break // connection already closed by client
		} else {
			fmt.Println("receive", string(request))
			temp := gocommand.DeCode(string(request))
			if temp == " !heart beat! " || temp == "" {
				//heart beat
			} else {
				receiveChan <- temp
			}
		}
		request = make([]byte, 1024) // clear last read content
	}
}