github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/sign/auto_sign_transport.go (about)

     1  package sign
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  
     7  	"github.com/artisanhe/tools/courier/transport_http"
     8  )
     9  
    10  var _ interface {
    11  	http.RoundTripper
    12  } = (*AutoSignTransport)(nil)
    13  
    14  func NewAutoSignTransport(secretExchanger SecretExchanger) transport_http.TransportWrapper {
    15  	return func(rt http.RoundTripper) http.RoundTripper {
    16  		if httpRt, ok := rt.(*http.Transport); ok {
    17  			autoSignTransport := &AutoSignTransport{
    18  				SecretExchanger: secretExchanger,
    19  				Transport:       httpRt,
    20  			}
    21  			return autoSignTransport
    22  		}
    23  		return rt
    24  	}
    25  }
    26  
    27  type AutoSignTransport struct {
    28  	SecretExchanger SecretExchanger
    29  	*http.Transport
    30  }
    31  
    32  func (t *AutoSignTransport) RoundTrip(req *http.Request) (*http.Response, error) {
    33  	if t.SecretExchanger == nil {
    34  		return nil, fmt.Errorf("missing SecretExchanger")
    35  	}
    36  
    37  	query := req.URL.Query()
    38  
    39  	sign, _, err := getSign(req, query, t.SecretExchanger)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  	if sign != nil {
    44  		query.Set(Sign, string(sign))
    45  		req.URL.RawQuery = query.Encode()
    46  	}
    47  
    48  	return t.Transport.RoundTrip(req)
    49  }