Golang TCPConn.CloseWrite方法代码示例

func tcpShim(inbound, outbound *net.TCPConn, connEvent *events.Connection, eh events.Handler) error {
	eh.Connection(connEvent)
	ch := make(chan error, 1)
	go func() {
		var err error
		defer func() { ch <- err }()
		_, err = io.Copy(inbound, outbound)
		outbound.CloseRead()
		inbound.CloseWrite()
	}()

	_, err1 := io.Copy(outbound, inbound)
	inbound.CloseRead()
	outbound.CloseWrite()

	err2 := <-ch
	inbound.Close()
	outbound.Close()

	if err1 != nil {
		return err1
	} else {
		return err2
	}
}

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

// 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 benchcore(host, req *string, port int, callback bench_callback) {
	var fail, bs, speed int64

	buf := make([]byte, 1500)
	reqLen := len(*req)
	var conn *net.TCPConn

	ds := fmt.Sprintf("%ds", *sec)

	dur, _ := time.ParseDuration(ds)
	timer := time.NewTimer(dur)

	for {
	nextconn:
		select {
		case <-timer.C:
			callback(fail, bs, speed)
			fmt.Println("Some of our childrens died.")
			return
		default:
			conn = socket.Socket(host, port)
			if conn == nil {
				fail++
				continue
			}
			wlen, err := conn.Write([]byte(*req))
			if wlen != reqLen || err != nil {
				fail++
				conn.Close()
				continue
			}
			if httpVersion == 0 {
				err := conn.CloseWrite()
				if err != nil {
					fail++
					conn.Close()
					continue
				}
			}
			if *force == false {
				for {
					b, err := conn.Read(buf)
					if err == io.EOF {
						break
					} else if err != nil {
						fail++
						conn.Close()
						break nextconn
					}
					bs += int64(b)
				}
			}
			if conn.Close() != nil {
				fail++
				continue
			}
			speed++
		}
	}
}

// Pipe starts bridging with two tcp connection
func Pipe(dst *net.TCPConn, src *net.TCPConn, f *func([]byte) []byte) error {
	defer src.CloseRead()
	defer dst.CloseWrite()

	rb := make([]byte, 4096)

	for {
		rsize, err := src.Read(rb)
		if err != nil {
			if isRecoverable(err) {
				continue
			}
			return err
		}

		var wb []byte
		if f != nil {
			wb = (*f)(rb[:rsize])
		} else {
			wb = rb[:rsize]
		}
		wWrote := 0
		wTotal := len(wb)
		for wWrote != wTotal {
			wSize, err := dst.Write(wb[wWrote:])
			wWrote += wSize
			if err != nil {
				if isRecoverable(err) {
					continue
				}
				return err
			}
		}
	}
}

func copyAndClose(ctx *ProxyCtx, dst, src *net.TCPConn) {
	if _, err := io.Copy(dst, src); err != nil {
		ctx.Warnf("Error copying to client: %s", err)
	}

	dst.CloseWrite()
	src.CloseRead()
}

func (this *VMessInboundHandler) HandleConnection(connection *net.TCPConn) error {
	defer connection.Close()

	connReader := v2net.NewTimeOutReader(16, connection)
	requestReader := protocol.NewVMessRequestReader(this.clients)

	request, err := requestReader.Read(connReader)
	if err != nil {
		log.Access(connection.RemoteAddr().String(), "", log.AccessRejected, err.Error())
		log.Warning("VMessIn: Invalid request from (%s): %v", connection.RemoteAddr().String(), err)
		return err
	}
	log.Access(connection.RemoteAddr().String(), request.Address.String(), log.AccessAccepted, "")
	log.Debug("VMessIn: Received request for %s", request.Address.String())

	ray := this.space.PacketDispatcher().DispatchToOutbound(v2net.NewPacket(request.Destination(), nil, true))
	input := ray.InboundInput()
	output := ray.InboundOutput()
	var readFinish, writeFinish sync.Mutex
	readFinish.Lock()
	writeFinish.Lock()

	userSettings := vmess.GetUserSettings(request.User.Level)
	connReader.SetTimeOut(userSettings.PayloadReadTimeout)
	go handleInput(request, connReader, input, &readFinish)

	responseKey := md5.Sum(request.RequestKey)
	responseIV := md5.Sum(request.RequestIV)

	aesStream, err := v2crypto.NewAesEncryptionStream(responseKey[:], responseIV[:])
	if err != nil {
		log.Error("VMessIn: Failed to create AES decryption stream: %v", err)
		close(input)
		return err
	}

	responseWriter := v2crypto.NewCryptionWriter(aesStream, connection)

	// Optimize for small response packet
	buffer := alloc.NewLargeBuffer().Clear()
	defer buffer.Release()
	buffer.AppendBytes(request.ResponseHeader[0] ^ request.ResponseHeader[1])
	buffer.AppendBytes(request.ResponseHeader[2] ^ request.ResponseHeader[3])
	buffer.AppendBytes(byte(0), byte(0))

	if data, open := <-output; open {
		buffer.Append(data.Value)
		data.Release()
		responseWriter.Write(buffer.Value)
		go handleOutput(request, responseWriter, output, &writeFinish)
		writeFinish.Lock()
	}

	connection.CloseWrite()
	readFinish.Lock()

	return nil
}

func sendSource(dest []*DestList, source *net.TCPConn) {
	for _, d := range dest {
		if !d.Discard {
			io.Copy(source, d.Conn)
		}
		d.Conn.CloseRead()
	}
	source.CloseWrite()
}

func copyBytes(in, out *net.TCPConn) {
	glog.Infof("Copying from %v <-> %v <-> %v <-> %v",
		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
	if _, err := io.Copy(in, out); err != nil {
		glog.Errorf("I/O error: %v", err)
	}
	in.CloseRead()
	out.CloseWrite()
}

func (self *Session) ProxyTCP(conn *net.TCPConn, bufferSize int) {
	writeClosed := make(chan struct{})
	readClosed := make(chan struct{})
	go func() {
		<-writeClosed
		<-readClosed
		conn.Close()
		self.log("proxy closed")
	}()
	// to conn
	go func() {
		var once sync.Once
		for {
			select {
			case msg := <-self.Message:
				switch msg.Tag {
				case DATA:
					_, err := conn.Write(msg.Data)
					if err != nil {
						self.log("proxy write error %v", err)
						go once.Do(func() {
							<-time.After(time.Second * 5)
							conn.CloseWrite()
							close(writeClosed)
						})
						self.AbortRead()
					}
				case STATE:
					switch msg.State {
					case STATE_FINISH_SEND, STATE_ABORT_SEND:
						go once.Do(func() {
							<-time.After(time.Second * 5)
							conn.CloseWrite()
							close(writeClosed)
						})
					case STATE_ABORT_READ:
					case STATE_FINISH_READ:
					}
				}
			case <-self.Stopped:
				return
			}
		}
	}()
	// from conn
	for {
		buf := make([]byte, bufferSize)
		n, err := conn.Read(buf)
		if err != nil {
			conn.CloseRead()
			close(readClosed)
			self.FinishSend()
			return
		}
		self.Send(buf[:n])
	}
}

// Copy one side of the socket, doing a half close when it has
// finished
func copy_half(backend *Backend, dst, src *net.TCPConn, wg *sync.WaitGroup) {
	defer wg.Done()
	transferred, err := io.Copy(dst, src)
	backend.transferred += transferred
	if err != nil {
		log.Printf("Error: %s", err)
	}
	dst.CloseWrite()
	src.CloseRead()
}

func copyConn(dst, src *net.TCPConn) {
	addConnection(src)
	_, err := io.Copy(dst, src)
	if err != nil {
		fmt.Println(err)
	}
	src.Close()
	dst.CloseWrite()
	removeConnection(src)
}

func handleConn(c *net.TCPConn) {
	input := bufio.NewScanner(c)
	for input.Scan() {
		wg.Add(1)
		go echo(c, input.Text(), 1*time.Second)
	}
	wg.Wait()
	c.CloseWrite()
	log.Println("close write")
}

func copyBytes(direction string, dest, src *net.TCPConn, wg *sync.WaitGroup) {
	defer wg.Done()
	glog.V(4).Infof("Copying %s: %s -> %s", direction, src.RemoteAddr(), dest.RemoteAddr())
	n, err := io.Copy(dest, src)
	if err != nil {
		glog.Errorf("I/O error: %v", err)
	}
	glog.V(4).Infof("Copied %d bytes %s: %s -> %s", n, direction, src.RemoteAddr(), dest.RemoteAddr())
	dest.CloseWrite()
	src.CloseRead()
}

func CopyBytes(in, out *net.TCPConn) {
	log.Printf("Copying from %v <-> %v <-> %v <-> %v",
		in.RemoteAddr(), in.LocalAddr(), out.LocalAddr(), out.RemoteAddr())
	_, err := io.Copy(in, out)
	if err != nil && err != io.EOF {
		log.Printf("I/O error: %v", err)
	}

	in.CloseRead()
	out.CloseWrite()
}