Golang io.ReadWriter類代碼示例

// dnsRoundTripUDP implements the dnsRoundTrip interface for RFC 1035's
// "UDP usage" transport mechanism. c should be a packet-oriented connection,
// such as a *UDPConn.
func dnsRoundTripUDP(c io.ReadWriter, query *dnsMsg) (*dnsMsg, error) {
	b, ok := query.Pack()
	if !ok {
		return nil, errors.New("cannot marshal DNS message")
	}
	if _, err := c.Write(b); err != nil {
		return nil, err
	}

	b = make([]byte, 512) // see RFC 1035
	for {
		n, err := c.Read(b)
		if err != nil {
			return nil, err
		}
		resp := &dnsMsg{}
		if !resp.Unpack(b[:n]) || !resp.IsResponseTo(query) {
			// Ignore invalid responses as they may be malicious
			// forgery attempts. Instead continue waiting until
			// timeout. See golang.org/issue/13281.
			continue
		}
		return resp, nil
	}
}

func negotiateMethod(c Context, rw io.ReadWriter) {
	rw.Write(methodRequest(proxy.MethodNoPassword))
	if _, err := rw.Read(make([]byte, 3)); err != nil {
		c.Error(err)
	}
	return
}

// dnsRoundTripTCP implements the dnsRoundTrip interface for RFC 1035's
// "TCP usage" transport mechanism. c should be a stream-oriented connection,
// such as a *TCPConn.
func dnsRoundTripTCP(c io.ReadWriter, query *dnsMsg) (*dnsMsg, error) {
	b, ok := query.Pack()
	if !ok {
		return nil, errors.New("cannot marshal DNS message")
	}
	l := len(b)
	b = append([]byte{byte(l >> 8), byte(l)}, b...)
	if _, err := c.Write(b); err != nil {
		return nil, err
	}

	b = make([]byte, 1280) // 1280 is a reasonable initial size for IP over Ethernet, see RFC 4035
	if _, err := io.ReadFull(c, b[:2]); err != nil {
		return nil, err
	}
	l = int(b[0])<<8 | int(b[1])
	if l > len(b) {
		b = make([]byte, l)
	}
	n, err := io.ReadFull(c, b[:l])
	if err != nil {
		return nil, err
	}
	resp := &dnsMsg{}
	if !resp.Unpack(b[:n]) {
		return nil, errors.New("cannot unmarshal DNS message")
	}
	if !resp.IsResponseTo(query) {
		return nil, errors.New("invalid DNS response")
	}
	return resp, nil
}

func (r *Receiver) handleTransfer(conn io.ReadWriter) (
	transfer.TransferResults, error,
) {
	if _, err := conn.Write([]byte("ok")); err != nil {
		return transfer.TransferResults{}, err
	}

	res := transfer.TransferResults{}
	buffer := make([]byte, 1024)

	startTime := time.Now()
	for {
		n, err := conn.Read(buffer)
		if err != nil { // done reading
			break
		}

		res.BytesSent += uint32(n)
		res.Checksum = crc32.Update(res.Checksum, crc32.IEEETable, buffer)
	}
	endTime := time.Now()

	res.Duration = endTime.Sub(startTime)

	return res, nil
}

// receiverEncHandshake negotiates a session token on conn.
// it should be called on the listening side of the connection.
//
// prv is the local client's private key.
// token is the token from a previous session with this node.
func receiverEncHandshake(conn io.ReadWriter, prv *ecdsa.PrivateKey, token []byte) (s secrets, err error) {
	authMsg := new(authMsgV4)
	authPacket, err := readHandshakeMsg(authMsg, encAuthMsgLen, prv, conn)
	if err != nil {
		return s, err
	}
	h := new(encHandshake)
	if err := h.handleAuthMsg(authMsg, prv); err != nil {
		return s, err
	}

	authRespMsg, err := h.makeAuthResp()
	if err != nil {
		return s, err
	}
	var authRespPacket []byte
	if authMsg.gotPlain {
		authRespPacket, err = authRespMsg.sealPlain(h)
	} else {
		authRespPacket, err = sealEIP8(authRespMsg, h)
	}
	if err != nil {
		return s, err
	}
	if _, err = conn.Write(authRespPacket); err != nil {
		return s, err
	}
	return h.secrets(authPacket, authRespPacket)
}

//throwBytes chucks bytes at the remote server then listens for a response
func throwBytes(conn io.ReadWriter, count uint64) error {
	var writeBytes uint64
	var b []byte
	buff := make([]byte, 128)
	for writeBytes < count {
		if (count - writeBytes) >= uint64(len(dataBlock)) {
			b = dataBlock
		} else {
			b = dataBlock[0:(count - writeBytes)]
		}
		n, err := conn.Write(b)
		if err != nil {
			return err
		}
		writeBytes += uint64(n)
	}
	//read the response
	n, err := conn.Read(buff)
	if err != nil {
		return err
	}
	if n == 0 {
		return fmt.Errorf("Failed to get OK on upload")
	}
	if !strings.HasPrefix(string(buff[0:n]), "OK ") {
		return fmt.Errorf("Failed to get OK on upload")
	}
	return nil
}

func handShake(rw io.ReadWriter) error {
	// version(1)+numMethods(1)+[256]methods
	buff := make([]byte, 258)
	n, err := io.ReadAtLeast(rw, buff, 2)
	if err != nil {
		return err
	}

	// version
	if buff[0] != 5 {
		return errors.New("socks unsuported version")
	}

	// numMethods
	numMethod := int(buff[1])

	// methods
	numMethod += 2
	if n <= numMethod {
		if _, err = io.ReadFull(rw, buff[n:numMethod]); err != nil {
			return err
		}
	}
	fmt.Println("Socks5: Step 1 -> ", buff[:numMethod])

	// return data
	buff[1] = 0 // no authentication
	if _, err := rw.Write(buff[:2]); err != nil {
		return err
	}
	fmt.Println("Socks5: Step 1 <- ", buff[:2])

	return nil
}

// swapKeys generates a public/private key pair and swaps the public key with
// a corresponding call to swapKeys over the given ReadWriter.
//
// swapKeys returns the private key generated locally and the public key of it's
// counterpart.
func swapKeys(rw io.ReadWriter) (priv, peer *[32]byte, err error) {
	// Always return nil for the keys if there is an error
	defer func() {
		if err != nil {
			priv, peer = nil, nil
		}
	}()
	pub, priv, err := box.GenerateKey(rand.Reader)
	if err != nil {
		return
	}

	// Write our public key
	werr := make(chan error)
	go func() {
		_, err := rw.Write(pub[:])
		werr <- err
	}()
	defer func() {
		if err == nil {
			err = <-werr
		}
	}()

	// Read their public key
	peer = new([32]byte)
	_, err = io.ReadFull(rw, peer[:])
	return
}

func performConnect(backend string, frontconn io.ReadWriter) {
	log.Printf("trying to connect to %s...\n", backend)
	backconn, err := net.Dial("tcp", backend)
	if err != nil {
		log.Printf("failed to connect to %s: %s\n", backend, err)
		frontconn.Write(errorReplyConnect(0x05))
		return
	}

	backaddr := backconn.RemoteAddr().String()
	log.Println("CONNECTED backend", backconn, backaddr)
	defer func() {
		backconn.Close()
		log.Println("DISCONNECTED backend", backconn, backaddr)
	}()

	// reply to the CONNECT command
	buf := make([]byte, 10)
	copy(buf, []byte{0x05, 0x00, 0x00, 0x01})
	packNetAddr(backconn.RemoteAddr(), buf[4:])
	frontconn.Write(buf)

	// bridge connection
	shutdown := make(chan bool, 2)
	go common.IOBridge(frontconn, backconn, shutdown)
	go common.IOBridge(backconn, frontconn, shutdown)

	<-shutdown
}

// Query writes a question to rw and waits for an answer.  It will pass the
// answer into the verify function.  Verify should check the answer for
// validity, returning a failure reason as a string, or an empty string if the
// answer is valid.
func QueryVerify(
	rw io.ReadWriter,
	question []byte,
	verify func(string) (string, error),
) (answer string, err error) {
	scanner := bufio.NewScanner(rw)
	for {
		_, err = rw.Write(question)
		if err != nil {
			return "", err
		}

		if !scanner.Scan() {
			if err = scanner.Err(); err != nil {
				return "", err
			}
			return "", fmt.Errorf("Connection closed")
		}
		answer = scanner.Text()
		failure, err := verify(answer)
		if err != nil {
			return "", err
		}
		if failure == "" {
			return answer, nil
		}
		_, err = fmt.Fprintln(rw, failure)
		if err != nil {
			return "", err
		}
	}
}

// initiatorEncHandshake negotiates a session token on conn.
// it should be called on the dialing side of the connection.
//
// prv is the local client's private key.
func initiatorEncHandshake(conn io.ReadWriter, prv *ecdsa.PrivateKey, remoteID discover.NodeID, token []byte) (s secrets, err error) {
	h := &encHandshake{initiator: true, remoteID: remoteID}
	authMsg, err := h.makeAuthMsg(prv, token)
	if err != nil {
		return s, err
	}
	var authPacket []byte
	if configSendEIP {
		authPacket, err = sealEIP8(authMsg, h)
	} else {
		authPacket, err = authMsg.sealPlain(h)
	}
	if err != nil {
		return s, err
	}
	if _, err = conn.Write(authPacket); err != nil {
		return s, err
	}

	authRespMsg := new(authRespV4)
	authRespPacket, err := readHandshakeMsg(authRespMsg, encAuthRespLen, prv, conn)
	if err != nil {
		return s, err
	}
	if err := h.handleAuthResp(authRespMsg); err != nil {
		return s, err
	}
	return h.secrets(authPacket, authRespPacket)
}

// ServeAgent serves the agent protocol on the given connection. It
// returns when an I/O error occurs.
func ServeAgent(agent Agent, c io.ReadWriter) error {
	s := &server{agent}

	var length [4]byte
	for {
		if _, err := io.ReadFull(c, length[:]); err != nil {
			return err
		}
		l := binary.BigEndian.Uint32(length[:])
		if l > maxAgentResponseBytes {
			// We also cap requests.
			return fmt.Errorf("agent: request too large: %d", l)
		}

		req := make([]byte, l)
		if _, err := io.ReadFull(c, req); err != nil {
			return err
		}

		repData := s.processRequestBytes(req)
		if len(repData) > maxAgentResponseBytes {
			return fmt.Errorf("agent: reply too large: %d bytes", len(repData))
		}

		binary.BigEndian.PutUint32(length[:], uint32(len(repData)))
		if _, err := c.Write(length[:]); err != nil {
			return err
		}
		if _, err := c.Write(repData); err != nil {
			return err
		}
	}
}

//從io接口中讀取指定結束符數據
func WriteAndReadWithTer(rw io.ReadWriter, writedata []byte, ter []byte, timeout int) ([]byte, error) {
	p := []byte(writedata)
	_, err := rw.Write(p)
	if err != nil {
		return nil, err
	}
	return ReadWithTer(rw, ter, timeout)
}

func (r *Receiver) handleBusy(conn io.ReadWriter) error {
	r.logger.Debugf("Server is busy!")
	if _, err := conn.Write([]byte("i-am-busy")); err != nil {
		return err
	}

	return nil
}

//從io接口中讀取指定長度數據
func WriteAndReadLen(rw io.ReadWriter, writedata []byte, readlen int, timeout int) ([]byte, error) {
	p := []byte(writedata)
	_, err := rw.Write(p)
	if err != nil {
		return nil, err
	}
	return ReadWithLen(rw, readlen, timeout)
}