trpc.group/trpc-go/trpc-go@v1.0.3/internal/attachment/attachment.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package attachment provides a internal implementation of tRPC client/server attachment.
    15  package attachment
    16  
    17  import (
    18  	"bytes"
    19  	"io"
    20  
    21  	"trpc.group/trpc-go/trpc-go/codec"
    22  )
    23  
    24  // ClientAttachmentKey is the key of client's Attachment.
    25  type ClientAttachmentKey struct{}
    26  
    27  // ServerAttachmentKey is the key of server's Attachment.
    28  type ServerAttachmentKey struct{}
    29  
    30  // Attachment stores the attachment in tRPC requests/responses.
    31  type Attachment struct {
    32  	Request  io.Reader
    33  	Response io.Reader
    34  }
    35  
    36  // NoopAttachment is an empty attachment.
    37  type NoopAttachment struct{}
    38  
    39  // Read implements the io.Reader interface, which always returns (0, io.EOF)
    40  func (a NoopAttachment) Read(_ []byte) (n int, err error) {
    41  	return 0, io.EOF
    42  }
    43  
    44  // ClientRequestAttachment returns client's Request Attachment from msg.
    45  func ClientRequestAttachment(msg codec.Msg) (io.Reader, bool) {
    46  	if a, _ := msg.CommonMeta()[ClientAttachmentKey{}].(*Attachment); a != nil {
    47  		return a.Request, true
    48  	}
    49  	return nil, false
    50  }
    51  
    52  // ServerResponseAttachment returns server's Response Attachment from msg.
    53  func ServerResponseAttachment(msg codec.Msg) (io.Reader, bool) {
    54  	if a, _ := msg.CommonMeta()[ServerAttachmentKey{}].(*Attachment); a != nil {
    55  		return a.Response, true
    56  	}
    57  	return nil, false
    58  }
    59  
    60  // SetClientResponseAttachment sets client's Response attachment to msg.
    61  // If the message does not contain client.Attachment,
    62  // which means that the user has explicitly ignored the att returned by the server.
    63  // For performance reasons, there is no need to set the response attachment into msg.
    64  func SetClientResponseAttachment(msg codec.Msg, attachment []byte) {
    65  	if a, _ := msg.CommonMeta()[ClientAttachmentKey{}].(*Attachment); a != nil {
    66  		a.Response = bytes.NewReader(attachment)
    67  	}
    68  }
    69  
    70  // SetServerRequestAttachment sets server's Request Attachment to msg.
    71  func SetServerRequestAttachment(m codec.Msg, attachment []byte) {
    72  	cm := m.CommonMeta()
    73  	if cm == nil {
    74  		cm = make(codec.CommonMeta)
    75  		m.WithCommonMeta(cm)
    76  	}
    77  	cm[ServerAttachmentKey{}] = &Attachment{Request: bytes.NewReader(attachment), Response: NoopAttachment{}}
    78  }