Golang buf.Buffer类代码示例

func (v *VMessOutboundHandler) handleRequest(session *encoding.ClientSession, conn internet.Connection, request *protocol.RequestHeader, payload *buf.Buffer, input buf.Reader, finish *sync.Mutex) {
	defer finish.Unlock()

	writer := bufio.NewWriter(conn)
	defer writer.Release()
	session.EncodeRequestHeader(request, writer)

	bodyWriter := session.EncodeRequestBody(request, writer)
	defer bodyWriter.Release()

	if !payload.IsEmpty() {
		if err := bodyWriter.Write(payload); err != nil {
			log.Info("VMess|Outbound: Failed to write payload. Disabling connection reuse.", err)
			conn.SetReusable(false)
		}
		payload.Release()
	}
	writer.SetCached(false)

	if err := buf.PipeUntilEOF(input, bodyWriter); err != nil {
		conn.SetReusable(false)
	}

	if request.Option.Has(protocol.RequestOptionChunkStream) {
		err := bodyWriter.Write(buf.NewLocal(8))
		if err != nil {
			conn.SetReusable(false)
		}
	}
	return
}

func (v *Stream) Write(data *buf.Buffer) (err error) {
	if data.IsEmpty() {
		return
	}

	select {
	case <-v.ctx.Done():
		return io.ErrClosedPipe
	case <-v.err:
		return io.ErrClosedPipe
	case <-v.close:
		return io.ErrClosedPipe
	default:
		select {
		case <-v.ctx.Done():
			return io.ErrClosedPipe
		case <-v.err:
			return io.ErrClosedPipe
		case <-v.close:
			return io.ErrClosedPipe
		case v.buffer <- data:
			v.inspector.Input(data)
			return nil
		}
	}
}

func (v *BlackHole) Dispatch(destination v2net.Destination, payload *buf.Buffer, ray ray.OutboundRay) {
	payload.Release()

	v.response.WriteTo(ray.OutboundOutput())
	ray.OutboundOutput().Close()

	ray.OutboundInput().Release()
}

func (v *DokodemoDoor) handleUDPResponse(dest v2net.Destination, payload *buf.Buffer) {
	defer payload.Release()
	v.udpMutex.RLock()
	defer v.udpMutex.RUnlock()
	if !v.accepting {
		return
	}
	v.udpHub.WriteTo(payload.Bytes(), dest)
}

func (w *UDPWriter) Write(b *buf.Buffer) error {
	eb := EncodeUDPPacket(w.request, b.Bytes())
	b.Release()
	defer eb.Release()
	if _, err := w.writer.Write(eb.Bytes()); err != nil {
		return err
	}
	return nil
}

func (v *ChunkWriter) Write(payload *buf.Buffer) error {
	totalLength := payload.Len()
	serial.Uint16ToBytes(uint16(totalLength), v.buffer[:0])
	v.auth.Authenticate(payload.Bytes())(v.buffer[2:])
	copy(v.buffer[2+AuthSize:], payload.Bytes())
	_, err := v.writer.Write(v.buffer[:2+AuthSize+payload.Len()])
	return err
}

func EncodeUDPPacket(request *protocol.RequestHeader, payload *buf.Buffer) (*buf.Buffer, error) {
	user := request.User
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, errors.Base(err).Message("Shadowsocks|UDP: Failed to parse account.")
	}
	account := rawAccount.(*ShadowsocksAccount)

	buffer := buf.NewSmall()
	ivLen := account.Cipher.IVSize()
	buffer.AppendSupplier(buf.ReadFullFrom(rand.Reader, ivLen))
	iv := buffer.Bytes()

	switch request.Address.Family() {
	case v2net.AddressFamilyIPv4:
		buffer.AppendBytes(AddrTypeIPv4)
		buffer.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyIPv6:
		buffer.AppendBytes(AddrTypeIPv6)
		buffer.Append([]byte(request.Address.IP()))
	case v2net.AddressFamilyDomain:
		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
		buffer.Append([]byte(request.Address.Domain()))
	default:
		return nil, errors.New("Shadowsocks|UDP: Unsupported address type: ", request.Address.Family())
	}

	buffer.AppendSupplier(serial.WriteUint16(uint16(request.Port)))
	buffer.Append(payload.Bytes())

	if request.Option.Has(RequestOptionOneTimeAuth) {
		authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
		buffer.SetByte(ivLen, buffer.Byte(ivLen)|0x10)

		buffer.AppendSupplier(authenticator.Authenticate(buffer.BytesFrom(ivLen)))
	}

	stream, err := account.Cipher.NewEncodingStream(account.Key, iv)
	if err != nil {
		return nil, errors.Base(err).Message("Shadowsocks|TCP: Failed to create encoding stream.")
	}

	stream.XORKeyStream(buffer.BytesFrom(ivLen), buffer.BytesFrom(ivLen))
	return buffer, nil
}

// Private: Visible for testing.
func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err != nil {
		log.Warning("DNS: Failed to parse DNS response: ", err)
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := DefaultTTL
	log.Debug("DNS: Handling response for id ", id, " content: ", msg.String())

	v.Lock()
	request, found := v.requests[id]
	if !found {
		v.Unlock()
		return
	}
	delete(v.requests, id)
	v.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (w *udpWorker) callback(b *buf.Buffer, source v2net.Destination, originalDest v2net.Destination) {
	conn, existing := w.getConnection(source)
	conn.input <- b.Bytes()

	if !existing {
		go func() {
			ctx := w.ctx
			ctx, cancel := context.WithCancel(ctx)
			conn.cancel = cancel
			if originalDest.IsValid() {
				ctx = proxy.ContextWithOriginalDestination(ctx, originalDest)
			}
			if len(w.tag) > 0 {
				ctx = proxy.ContextWithInboundTag(ctx, w.tag)
			}
			ctx = proxy.ContextWithSource(ctx, source)
			ctx = proxy.ContextWithInboundDestination(ctx, v2net.UDPDestination(w.address, w.port))
			w.proxy.Process(ctx, v2net.Network_UDP, conn)
			w.removeConn(source)
			cancel()
		}()
	}
}

func (v *Server) handlerUDPPayload(payload *buf.Buffer, session *proxy.SessionInfo) {
	source := session.Source
	request, data, err := DecodeUDPPacket(v.user, payload)
	if err != nil {
		log.Info("Shadowsocks|Server: Skipping invalid UDP packet from: ", source, ": ", err)
		log.Access(source, "", log.AccessRejected, err)
		payload.Release()
		return
	}

	if request.Option.Has(RequestOptionOneTimeAuth) && v.account.OneTimeAuth == Account_Disabled {
		log.Info("Shadowsocks|Server: Client payload enables OTA but server doesn't allow it.")
		payload.Release()
		return
	}

	if !request.Option.Has(RequestOptionOneTimeAuth) && v.account.OneTimeAuth == Account_Enabled {
		log.Info("Shadowsocks|Server: Client payload disables OTA but server forces it.")
		payload.Release()
		return
	}

	dest := request.Destination()
	log.Access(source, dest, log.AccessAccepted, "")
	log.Info("Shadowsocks|Server: Tunnelling request to ", dest)

	v.udpServer.Dispatch(&proxy.SessionInfo{Source: source, Destination: dest, User: request.User, Inbound: v.meta}, data, func(destination v2net.Destination, payload *buf.Buffer) {
		defer payload.Release()

		data, err := EncodeUDPPacket(request, payload)
		if err != nil {
			log.Warning("Shadowsocks|Server: Failed to encode UDP packet: ", err)
			return
		}
		defer data.Release()

		v.udpHub.WriteTo(data.Bytes(), source)
	})
}

func (v *Server) handleUDPPayload(payload *buf.Buffer, session *proxy.SessionInfo) {
	source := session.Source
	log.Info("Socks: Client UDP connection from ", source)
	request, err := protocol.ReadUDPRequest(payload.Bytes())
	payload.Release()

	if err != nil {
		log.Error("Socks: Failed to parse UDP request: ", err)
		return
	}
	if request.Data.Len() == 0 {
		request.Data.Release()
		return
	}
	if request.Fragment != 0 {
		log.Warning("Socks: Dropping fragmented UDP packets.")
		// TODO handle fragments
		request.Data.Release()
		return
	}

	log.Info("Socks: Send packet to ", request.Destination(), " with ", request.Data.Len(), " bytes")
	log.Access(source, request.Destination, log.AccessAccepted, "")
	v.udpServer.Dispatch(&proxy.SessionInfo{Source: source, Destination: request.Destination(), Inbound: v.meta}, request.Data, func(destination v2net.Destination, payload *buf.Buffer) {
		response := &protocol.Socks5UDPRequest{
			Fragment: 0,
			Address:  request.Destination().Address,
			Port:     request.Destination().Port,
			Data:     payload,
		}
		log.Info("Socks: Writing back UDP response with ", payload.Len(), " bytes to ", destination)

		udpMessage := buf.NewLocal(2048)
		response.Write(udpMessage)

		v.udpMutex.RLock()
		if !v.accepting {
			v.udpMutex.RUnlock()
			return
		}
		nBytes, err := v.udpHub.WriteTo(udpMessage.Bytes(), destination)
		v.udpMutex.RUnlock()
		udpMessage.Release()
		response.Data.Release()
		if err != nil {
			log.Error("Socks: failed to write UDP message (", nBytes, " bytes) to ", destination, ": ", err)
		}
	})
}

func DecodeUDPPacket(user *protocol.User, payload *buf.Buffer) (*protocol.RequestHeader, *buf.Buffer, error) {
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		return nil, nil, errors.Base(err).Message("Shadowsocks|UDP: Failed to parse account.")
	}
	account := rawAccount.(*ShadowsocksAccount)

	ivLen := account.Cipher.IVSize()
	iv := payload.BytesTo(ivLen)
	payload.SliceFrom(ivLen)

	stream, err := account.Cipher.NewDecodingStream(account.Key, iv)
	if err != nil {
		return nil, nil, errors.Base(err).Message("Shadowsocks|UDP: Failed to initialize decoding stream.")
	}
	stream.XORKeyStream(payload.Bytes(), payload.Bytes())

	authenticator := NewAuthenticator(HeaderKeyGenerator(account.Key, iv))
	request := &protocol.RequestHeader{
		Version: Version,
		User:    user,
		Command: protocol.RequestCommandUDP,
	}

	addrType := (payload.Byte(0) & 0x0F)
	if (payload.Byte(0) & 0x10) == 0x10 {
		request.Option |= RequestOptionOneTimeAuth
	}

	if request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Disabled {
		return nil, nil, errors.New("Shadowsocks|UDP: Rejecting packet with OTA enabled, while server disables OTA.")
	}

	if !request.Option.Has(RequestOptionOneTimeAuth) && account.OneTimeAuth == Account_Enabled {
		return nil, nil, errors.New("Shadowsocks|UDP: Rejecting packet with OTA disabled, while server enables OTA.")
	}

	if request.Option.Has(RequestOptionOneTimeAuth) {
		payloadLen := payload.Len() - AuthSize
		authBytes := payload.BytesFrom(payloadLen)

		actualAuth := make([]byte, AuthSize)
		authenticator.Authenticate(payload.BytesTo(payloadLen))(actualAuth)
		if !bytes.Equal(actualAuth, authBytes) {
			return nil, nil, errors.New("Shadowsocks|UDP: Invalid OTA.")
		}

		payload.Slice(0, payloadLen)
	}

	payload.SliceFrom(1)

	switch addrType {
	case AddrTypeIPv4:
		request.Address = v2net.IPAddress(payload.BytesTo(4))
		payload.SliceFrom(4)
	case AddrTypeIPv6:
		request.Address = v2net.IPAddress(payload.BytesTo(16))
		payload.SliceFrom(16)
	case AddrTypeDomain:
		domainLength := int(payload.Byte(0))
		request.Address = v2net.DomainAddress(string(payload.BytesRange(1, 1+domainLength)))
		payload.SliceFrom(1 + domainLength)
	default:
		return nil, nil, errors.New("Shadowsocks|UDP: Unknown address type: ", addrType)
	}

	request.Port = v2net.PortFromBytes(payload.BytesTo(2))
	payload.SliceFrom(2)

	return request, payload, nil
}

func (v *Listener) OnReceive(payload *buf.Buffer, src v2net.Destination, originalDest v2net.Destination) {
	defer payload.Release()

	segments := v.reader.Read(payload.Bytes())
	if len(segments) == 0 {
		log.Info("KCP|Listener: discarding invalid payload from ", src)
		return
	}

	if !v.running {
		return
	}
	v.Lock()
	defer v.Unlock()
	if !v.running {
		return
	}
	if payload.Len() < 4 {
		return
	}
	conv := segments[0].Conversation()
	cmd := segments[0].Command()

	id := ConnectionID{
		Remote: src.Address,
		Port:   src.Port,
		Conv:   conv,
	}
	conn, found := v.sessions[id]

	if !found {
		if cmd == CommandTerminate {
			return
		}
		writer := &Writer{
			id:       id,
			hub:      v.hub,
			dest:     src,
			listener: v,
		}
		remoteAddr := &net.UDPAddr{
			IP:   src.Address.IP(),
			Port: int(src.Port),
		}
		localAddr := v.hub.Addr()
		sConn := &ServerConnection{
			id:     internal.NewConnectionID(v2net.LocalHostIP, src),
			local:  localAddr,
			remote: remoteAddr,
			writer: &KCPPacketWriter{
				Header:   v.header,
				Writer:   writer,
				Security: v.security,
			},
			closer: writer,
		}
		conn = NewConnection(conv, sConn, v, v.config)
		select {
		case v.awaitingConns <- conn:
		case <-time.After(time.Second * 5):
			conn.Close()
			return
		}
		v.sessions[id] = conn
	}
	conn.Input(segments)
}

func (request *Socks5UDPRequest) Write(buffer *buf.Buffer) {
	buffer.AppendBytes(0, 0, request.Fragment)
	switch request.Address.Family() {
	case v2net.AddressFamilyIPv4:
		buffer.AppendBytes(AddrTypeIPv4)
		buffer.Append(request.Address.IP())
	case v2net.AddressFamilyIPv6:
		buffer.AppendBytes(AddrTypeIPv6)
		buffer.Append(request.Address.IP())
	case v2net.AddressFamilyDomain:
		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
		buffer.Append([]byte(request.Address.Domain()))
	}
	buffer.AppendSupplier(serial.WriteUint16(request.Port.Value()))
	buffer.Append(request.Data.Bytes())
}

func (v *Client) Dispatch(destination v2net.Destination, payload *buf.Buffer, ray ray.OutboundRay) {
	defer payload.Release()
	defer ray.OutboundInput().Release()
	defer ray.OutboundOutput().Close()

	network := destination.Network

	var server *protocol.ServerSpec
	var conn internet.Connection

	err := retry.ExponentialBackoff(5, 100).On(func() error {
		server = v.serverPicker.PickServer()
		dest := server.Destination()
		dest.Network = network
		rawConn, err := internet.Dial(v.meta.Address, dest, v.meta.GetDialerOptions())
		if err != nil {
			return err
		}
		conn = rawConn

		return nil
	})
	if err != nil {
		log.Warning("Shadowsocks|Client: Failed to find an available destination:", err)
		return
	}
	log.Info("Shadowsocks|Client: Tunneling request to ", destination, " via ", server.Destination())

	conn.SetReusable(false)

	request := &protocol.RequestHeader{
		Version: Version,
		Address: destination.Address,
		Port:    destination.Port,
	}
	if destination.Network == v2net.Network_TCP {
		request.Command = protocol.RequestCommandTCP
	} else {
		request.Command = protocol.RequestCommandUDP
	}

	user := server.PickUser()
	rawAccount, err := user.GetTypedAccount()
	if err != nil {
		log.Warning("Shadowsocks|Client: Failed to get a valid user account: ", err)
		return
	}
	account := rawAccount.(*ShadowsocksAccount)
	request.User = user

	if account.OneTimeAuth == Account_Auto || account.OneTimeAuth == Account_Enabled {
		request.Option |= RequestOptionOneTimeAuth
	}

	if request.Command == protocol.RequestCommandTCP {
		bufferedWriter := bufio.NewWriter(conn)
		defer bufferedWriter.Release()

		bodyWriter, err := WriteTCPRequest(request, bufferedWriter)
		defer bodyWriter.Release()

		if err != nil {
			log.Info("Shadowsock|Client: Failed to write request: ", err)
			return
		}

		if err := bodyWriter.Write(payload); err != nil {
			log.Info("Shadowsocks|Client: Failed to write payload: ", err)
			return
		}

		var responseMutex sync.Mutex
		responseMutex.Lock()

		go func() {
			defer responseMutex.Unlock()

			responseReader, err := ReadTCPResponse(user, conn)
			if err != nil {
				log.Warning("Shadowsocks|Client: Failed to read response: ", err)
				return
			}

			if err := buf.PipeUntilEOF(responseReader, ray.OutboundOutput()); err != nil {
				log.Info("Shadowsocks|Client: Failed to transport all TCP response: ", err)
			}
		}()

		bufferedWriter.SetCached(false)
		if err := buf.PipeUntilEOF(ray.OutboundInput(), bodyWriter); err != nil {
			log.Info("Shadowsocks|Client: Failed to trasnport all TCP request: ", err)
		}

		responseMutex.Lock()
	}

	if request.Command == protocol.RequestCommandUDP {
		timedReader := v2net.NewTimeOutReader(16, conn)
		var responseMutex sync.Mutex
		responseMutex.Lock()
//.........这里部分代码省略.........