trpc.group/trpc-go/trpc-go@v1.0.3/server/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 server
    15  
    16  import (
    17  	"io"
    18  
    19  	"trpc.group/trpc-go/trpc-go/codec"
    20  	"trpc.group/trpc-go/trpc-go/internal/attachment"
    21  )
    22  
    23  // Attachment stores the attachment of tRPC requests/responses.
    24  type Attachment struct {
    25  	attachment *attachment.Attachment
    26  }
    27  
    28  // Request returns Request Attachment.
    29  func (a *Attachment) Request() io.Reader {
    30  	return a.attachment.Request
    31  }
    32  
    33  // SetResponse sets Response attachment.
    34  func (a *Attachment) SetResponse(attachment io.Reader) {
    35  	a.attachment.Response = attachment
    36  }
    37  
    38  // GetAttachment returns Attachment from msg.
    39  // If there is no Attachment in the msg, an empty attachment bound to the msg will be returned.
    40  func GetAttachment(msg codec.Msg) *Attachment {
    41  	cm := msg.CommonMeta()
    42  	if cm == nil {
    43  		cm = make(codec.CommonMeta)
    44  		msg.WithCommonMeta(cm)
    45  	}
    46  	a, _ := cm[attachment.ServerAttachmentKey{}]
    47  	if a == nil {
    48  		a = &attachment.Attachment{Request: attachment.NoopAttachment{}, Response: attachment.NoopAttachment{}}
    49  		cm[attachment.ServerAttachmentKey{}] = a
    50  	}
    51  
    52  	return &Attachment{attachment: a.(*attachment.Attachment)}
    53  }