Golang net.TCPConn类代码示例

//
// 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
}

// handler handles a single TCP Connection
func (t *TcpListener) handler(conn *net.TCPConn, id string) {
	// connection cleanup function
	defer func() {
		t.wg.Done()
		conn.Close()
		log.Printf("Closed TCP Connection from %s", conn.RemoteAddr())
		// Add one connection potential back to channel when this one closes
		t.accept <- true
		t.forget(id)
	}()

	scanner := bufio.NewScanner(conn)
	for {
		select {
		case <-t.done:
			return
		default:
			if !scanner.Scan() {
				return
			}
			buf := scanner.Bytes()
			select {
			case t.in <- buf:
			default:
				log.Printf(dropwarn, string(buf))
			}
		}
	}
}

func listenForMoveRequests(conn *net.TCPConn) {
	listenServe := make(chan []uint8)
	easynet.TieConnToChannel(conn, listenServe)
	for data := range listenServe {
		r := new(ttypes.BotMoveRequest)
		err := json.Unmarshal(data, r)
		easynet.DieIfError(err, "JSON error")

		if r.Kill {
			fmt.Printf("Bot on %s received kill signal\n", os.Args[0])
			os.Exit(0)
		}

		fmt.Printf("Bot at %d, %d received messages: %v\n", r.YourX, r.YourY, r.Messages)
		fmt.Printf("  Sees other bots: %v\n", r.OtherBots)
		//Do something

		response := new(ttypes.BotMoveResponse)
		if r.YourY < 5 {
			response.MoveDirection = "up"
		} else if r.YourY > 5 {
			response.MoveDirection = "down"
		}
		// response.BroadcastMessage = fmt.Sprintf("'I am %v at %d, %d'", os.Args[0], r.YourX, r.YourY)
		responseString, err := json.Marshal(response)
		easynet.DieIfError(err, "JSON marshal error")
		conn.Write(responseString)
	}
}

func (hs *SimpleHandshake) WithClient(conn *Conn) error {
	var iorw *net.TCPConn = conn.IoRw

	c0c1 := make([]byte, 1537)
	conn.Logger.Info("read c0c1 from conn, size=%v", len(c0c1))
	if _, err := io.ReadFull(iorw, c0c1); err != nil {
		conn.Logger.Error("read c0c1 failed, err is %v", err)
		return err
	}
	conn.Logger.Info("read c0c1 ok")

	if c0c1[0] != 0x03 {
		conn.Logger.Error("rtmp plain required 0x03, actual is %#x", c0c1[0])
		return RtmpPlainRequired
	}
	conn.Logger.Info("check rtmp plain protocol ok")

	// use bytes buffer to write content.
	s0s1s2 := bytes.NewBuffer(make([]byte, 0, 3073))

	// plain text required.
	binary.Write(s0s1s2, binary.BigEndian, byte(0x03))
	// s1 time
	binary.Write(s0s1s2, binary.BigEndian, int32(time.Now().Unix()))
	// s1 time2 copy from c1
	if _, err := s0s1s2.Write(c0c1[1:5]); err != nil {
		conn.Logger.Error("copy c0c1 time to s0s1s2 failed, err is %v", err)
		return err
	}
	// s1 1528 random bytes
	s0s1s2Random := make([]byte, 1528)
	RandomGenerate(conn.Rand, s0s1s2Random)
	if _, err := s0s1s2.Write(s0s1s2Random); err != nil {
		conn.Logger.Error("fill s1 random bytes failed, err is %v", err)
		return err
	}
	// if c1 specified, copy c1 to s2.
	// @see: https://github.com/winlinvip/simple-rtmp-server/issues/46
	if _, err := s0s1s2.Write(c0c1[1:1537]); err != nil {
		conn.Logger.Error("copy c1 to s1 failed, err is %v", err)
		return err
	}
	conn.Logger.Info("generate s0s1s2 ok, buf=%d", s0s1s2.Len())

	if written, err := iorw.Write(s0s1s2.Bytes()); err != nil {
		conn.Logger.Error("send s0s1s2 failed, written=%d, err is %v", written, err)
		return err
	}
	conn.Logger.Info("send s0s1s2 ok")

	c2 := make([]byte, 1536)
	conn.Logger.Info("read c2 from conn, size=%v", len(c2))
	if _, err := io.ReadFull(iorw, c2); err != nil {
		conn.Logger.Error("read c2 failed, err is %v", err)
		return err
	}
	conn.Logger.Info("read c2 ok")

	return nil
}

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
}

func newClientConn(ls *LeaseServer, raw *net.TCPConn) (*clientConn, error) {
	err := raw.SetKeepAlive(true)
	if err != nil {
		return nil, err
	}
	err = raw.SetNoDelay(true)
	if err != nil {
		return nil, err
	}
	ret := &clientConn{
		id:          incrementAndGet(&ls.clientIdCounter, 1),
		c:           raw,
		d:           gob.NewDecoder(raw),
		e:           gob.NewEncoder(raw),
		req:         ls.req,
		resp:        make(chan response, 20),
		ackLock:     new(sync.Mutex),
		pendingAcks: make(map[uint64]chan bool),
	}
	// send client id
	fmt.Printf("sending id %d\n", ret.id)
	idBuff := make([]byte, 8, 8)
	binary.LittleEndian.PutUint64(idBuff, ret.id)
	ret.c.Write(idBuff)
	fmt.Println("sent")
	return ret, nil
}

// Copy from WebSocket to socket and vice versa.
func proxy(local *net.TCPConn, conn *websocketConn) {
	var wg sync.WaitGroup

	wg.Add(2)

	go func() {
		_, err := io.Copy(conn, local)
		if err != nil {
			Log("error copying ORPort to WebSocket")
		}
		local.CloseRead()
		conn.Close()
		wg.Done()
	}()

	go func() {
		_, err := io.Copy(local, conn)
		if err != nil {
			Log("error copying WebSocket to ORPort")
		}
		local.CloseWrite()
		conn.Close()
		wg.Done()
	}()

	wg.Wait()
}

func (t *Listener) accept(c *net.TCPConn) {
	g := newGobConn(c)

	/*
		XXX: Maybe this handshake should be in auto, where the other side of it is
		dmsg_, err := g.Read()
		if err != nil {
			g.Close()
			return
		}
		dmsg, ok := dmsg_.(*autoDialMsg)
		if !ok {
			g.Close()
			return
		}
		if err := g.Write(&autoAcceptMsg{}); err != nil {
			g.Close()
			return
		}
	*/

	addr := x.Addr(c.RemoteAddr().String())
	t.olk.Lock()
	defer t.olk.Unlock()
	l := t.open[dmsg.ID]
	if l == nil {
		l = newAcceptLink(addr, dmsg.ID, g, listenerBroker{t})
		t.open[dmsg.ID] = l
	} else {
		l.AcceptRedial(g)
	}
}

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
		}
	}
}

// Read from socket and write to websocket
func ReadSocket(ws *websocket.Conn, conn *net.TCPConn) {
	reader := bufio.NewReader(conn)
	var line string = ""
	for {
		if reader == nil {
			break
		}
		buffer, isPrefix, err := reader.ReadLine()
		if err != nil {
			break
		}
		// fmt.Println("ReadSocket: got", len(buffer), "bytes")
		line = line + string(buffer)
		if !isPrefix {
			// fmt.Println("Sending message to web socket:", line)
			err = websocket.Message.Send(ws, line)
			if err != nil {
				_ = conn.Close()
			}
			line = ""
		}
	}
	fmt.Println("ReadSocket exit")
	ws.Close()
}

func handle(c *net.TCPConn, config Config) {
	defer c.Close()
	// TODO c.SetTimeout(60e9)
	r := bufio.NewReaderSize(c, 4096)
	for {

		// Note that everything in this loop should proceed as fast as it can
		// so we're not blocked and can keep processing
		// so the validation, the pipeline initiated via table.Dispatch(), etc
		// must never block.

		// note that we don't support lines longer than 4096B. that seems very reasonable..
		buf, _, err := r.ReadLine()

		if nil != err {
			if io.EOF != err {
				log.Error(err.Error())
			}
			break
		}

		buf_copy := make([]byte, len(buf), len(buf))
		copy(buf_copy, buf)
		numIn.Inc(1)

		table.Dispatch(buf_copy)
	}
}

func (consumer *BrokerConsumer) tryConnect(conn *net.TCPConn, tp *TopicPartition) (err error, reader *ByteBuffer) {
	var errCode int
	request := consumer.broker.EncodeConsumeRequest()
	//log.Println("offset=", tp.Offset, " ", tp.MaxSize, " ", tp.Topic, " ", tp.Partition, "  \n\t", string(request), request)
	_, err = conn.Write(request)
	if err != nil {
		if err = consumer.handleConnError(err, conn); err != nil {
			return err, nil
		}
	}

	reader = consumer.broker.readResponse(conn)
	err, errCode = reader.ReadHeader()
	if err != nil && errCode == 1 {
		log.Println("Bad Offset id, resetting?")
		// Error Code 1 means bad offsetid, we shold get a good offset, and reconnect!
		offsetVal := GetOffset(consumer.broker.hostname, tp)
		if offsetVal > 0 {
			// RECONNECT!
			log.Println("RECONNECTING !!! ", offsetVal)
			tp.Offset = offsetVal
			if err, reader = consumer.tryConnect(conn, tp); err != nil {
				return err, nil
			}
		} else {
			return err, nil
		}

	} else if err != nil {
		//log.Println("offset=", tp.Offset, " ", tp.MaxSize, " ", err.Error(), " ", request, " ", tp.Topic, " ", tp.Partition, "  \n\t", string(request))
		return err, nil
	}
	return
}

// NewClient returns new TCP client connection.
func NewClient(md *meta.Data, serverAddress string, sockerPort, providerPort int) (client *Client, err error) {
	var tcpAddr *net.TCPAddr
	tcpAddr, err = net.ResolveTCPAddr("tcp4", serverAddress)
	if err != nil {
		return
	}

	var tcpConn *net.TCPConn
	tcpConn, err = net.DialTCP("tcp", nil, tcpAddr)
	if err != nil {
		return
	}

	tcpConn.Write(append([]byte(md.UserID+"::"+strconv.Itoa(sockerPort)+"::"+strconv.Itoa(providerPort)), '\n'))

	client = &Client{
		TCPConn: tcpConn,
		close:   make(chan bool),
		ackChan: make(chan bool),
		reader:  bufio.NewReader(tcpConn),
		md:      md,
	}

	return
}

func (srvr *FileServer) serve(conn *net.TCPConn, clientID int) {
	reader := bufio.NewReader(conn)
	for {
		msg, msgerr, fatalerr := fs.GetMsg(reader)
		leader := srvr.rn.LeaderId()
		if leader != srvr.id {
			//fmt.Println(srvr.id, " : To client ID ", msg, " : Not the Leader, sending an L ")
			if leader == -1 {
				// Don't know the leader
				//fmt.Println("Do not know the current leader")
				reply(conn, &fs.Msg{Kind: 'L', Contents: []byte("ERR_REDIRECT _")})
			} else {
				content := []byte("ERR_REDIRECT " + srvr.nodeURLs[leader])
				//fmt.Println("Redirecting to correct URL : ",  srvr.nodeURLs[leader])
				reply(conn, &fs.Msg{Kind: 'L', Contents: content, Numbytes: len(content)})
			}
			break
		}

		if fatalerr != nil || msgerr != nil {
			reply(conn, &fs.Msg{Kind: 'M'})
			conn.Close()
			break
		}
		msg.ClientID = clientID
		msg.ServerID = srvr.id
		data, err := json.Marshal(msg)
		//fmt.Println("Marshalled data : ", string(data))
		check(err)
		srvr.rn.Append(data)

	}
}

func proxy(cliConn *net.TCPConn, rAddr *net.TCPAddr) error {
	srvConn, err := net.DialTCP("tcp", nil, rAddr)
	if err != nil {
		cliConn.Close()
		return err
	}
	defer srvConn.Close()

	// channels to wait on the close event for each connection
	serverClosed := make(chan struct{}, 1)
	clientClosed := make(chan struct{}, 1)

	go broker(srvConn, cliConn, clientClosed)
	go broker(cliConn, srvConn, serverClosed)

	var waitFor chan struct{}
	select {
	case <-clientClosed:
		// the client closed first
		srvConn.SetLinger(0)
		srvConn.CloseRead()
		waitFor = serverClosed
	case <-serverClosed:
		cliConn.CloseRead()
		waitFor = clientClosed
	}

	<-waitFor
	return nil
}