github.com/aporeto-inc/trireme-lib@v10.358.0+incompatible/controller/internal/enforcer/applicationproxy/http/transport.go (about)

     1  package httpproxy
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"go.aporeto.io/enforcerd/trireme-lib/controller/internal/enforcer/flowstats"
     7  )
     8  
     9  // TriremeRoundTripper is the Trireme RoundTripper that will handle
    10  // responses.
    11  type TriremeRoundTripper struct {
    12  	http.RoundTripper
    13  }
    14  
    15  // NewTriremeRoundTripper creates a new RoundTripper that handles the
    16  // responses.
    17  func NewTriremeRoundTripper(r http.RoundTripper) *TriremeRoundTripper {
    18  	return &TriremeRoundTripper{
    19  		RoundTripper: r,
    20  	}
    21  }
    22  
    23  // RoundTrip implements the RoundTripper interface. It will add a cookie
    24  // in the response in case of OIDC requests with refresh tokens.
    25  func (t *TriremeRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
    26  
    27  	res, err := t.RoundTripper.RoundTrip(req)
    28  	if err != nil || res == nil {
    29  		return res, err
    30  	}
    31  
    32  	data := req.Context().Value(statsContextKey)
    33  	if data == nil {
    34  		return res, nil
    35  	}
    36  
    37  	state, ok := data.(*flowstats.ConnectionState)
    38  	if ok && state.Cookie == nil {
    39  		return res, nil
    40  	}
    41  
    42  	if v := state.Cookie.String(); v != "" {
    43  		res.Header.Add("Set-Cookie", v)
    44  	}
    45  
    46  	return res, nil
    47  }