1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
|
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package flight
import (
"context"
"net/http"
"strings"
"sync"
"time"
"golang.org/x/exp/maps"
"google.golang.org/grpc/metadata"
)
// endOfTime is the time when session (non-persistent) cookies expire.
// This instant is representable in most date/time formats (not just
// Go's time.Time) and should be far enough in the future.
// taken from Go's net/http/cookiejar/jar.go
var endOfTime = time.Date(9999, 12, 31, 23, 59, 59, 0, time.UTC)
// NewClientCookieMiddleware returns a go-routine safe middleware for flight
// clients which properly handles Set-Cookie headers to store cookies
// in a cookie jar, and then requests are sent with those cookies added
// as a Cookie header.
func NewClientCookieMiddleware() ClientMiddleware {
return CreateClientMiddleware(&clientCookieMiddleware{jar: make(map[string]http.Cookie)})
}
func NewCookieMiddleware() CookieMiddleware {
return &clientCookieMiddleware{jar: make(map[string]http.Cookie)}
}
// CookieMiddleware is a go-routine safe middleware for flight clients
// which properly handles Set-Cookie headers for storing cookies.
// This can be passed into `CreateClientMiddleware` to create a new
// middleware object. You can also clone it to create middleware for a
// new client which starts with the same cookies.
type CookieMiddleware interface {
CustomClientMiddleware
// Clone creates a new CookieMiddleware that starts out with the same
// cookies that this one already has. This is useful when creating a
// new client connection for the same server.
Clone() CookieMiddleware
}
type clientCookieMiddleware struct {
jar map[string]http.Cookie
mx sync.Mutex
}
func (cc *clientCookieMiddleware) Clone() CookieMiddleware {
cc.mx.Lock()
defer cc.mx.Unlock()
return &clientCookieMiddleware{jar: maps.Clone(cc.jar)}
}
func (cc *clientCookieMiddleware) StartCall(ctx context.Context) context.Context {
cc.mx.Lock()
defer cc.mx.Unlock()
if len(cc.jar) == 0 {
return ctx
}
now := time.Now()
// Per RFC 6265 section 5.4, rather than adding multiple cookie strings
// or multiple cookie headers, multiple cookies are all sent as a single
// header value separated by semicolons.
// we will also clear any expired cookies from the jar while we determine
// the cookies to send.
cookies := make([]string, 0, len(cc.jar))
for id, c := range cc.jar {
if !c.Expires.After(now) {
delete(cc.jar, id)
continue
}
cookies = append(cookies, (&http.Cookie{Name: c.Name, Value: c.Value}).String())
}
if len(cookies) == 0 {
return ctx
}
return metadata.AppendToOutgoingContext(ctx, "Cookie", strings.Join(cookies, ";"))
}
func processCookieExpire(c *http.Cookie, now time.Time) (remove bool) {
// MaxAge takes precedence over Expires
if c.MaxAge < 0 {
return true
} else if c.MaxAge > 0 {
c.Expires = now.Add(time.Duration(c.MaxAge) * time.Second)
} else {
if c.Expires.IsZero() {
c.Expires = endOfTime
} else {
if !c.Expires.After(now) {
return true
}
}
}
return
}
func (cc *clientCookieMiddleware) HeadersReceived(ctx context.Context, md metadata.MD) {
// instead of replicating the logic for processing the Set-Cookie
// header, let's just make a fake response and use the built-in
// cookie processing. It's very non-trivial
cookies := (&http.Response{
Header: http.Header{"Set-Cookie": md.Get("set-cookie")},
}).Cookies()
now := time.Now()
cc.mx.Lock()
defer cc.mx.Unlock()
for _, c := range cookies {
id := c.Name + c.Path
if processCookieExpire(c, now) {
delete(cc.jar, id)
continue
}
cc.jar[id] = *c
}
}
|
