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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
|
package database
import (
"context"
"database/sql"
"errors"
"fmt"
"net/http"
"time"
"github.com/ansg191/ibd-trader/backend/internal/keys"
)
func GetAnyCookie(ctx context.Context, exec Executor, kms keys.KeyManagementService) (*IBDCookie, error) {
row := exec.QueryRowContext(ctx, `
SELECT ibd_tokens.id, token, encrypted_key, kms_key_name, expires_at
FROM ibd_tokens
INNER JOIN keys ON encryption_key = keys.id
WHERE expires_at > NOW()
AND degraded = FALSE
ORDER BY random()
LIMIT 1;`)
var id uint
var encryptedToken, encryptedKey []byte
var keyName string
var expiry time.Time
err := row.Scan(&id, &encryptedToken, &encryptedKey, &keyName, &expiry)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, nil
}
return nil, fmt.Errorf("unable to scan sql row into ibd cookie: %w", err)
}
// Set the expiry to UTC explicitly.
// For some reason, the expiry time is set to location="".
expiry = expiry.UTC()
token, err := keys.Decrypt(ctx, kms, keyName, encryptedToken, encryptedKey)
if err != nil {
return nil, fmt.Errorf("unable to decrypt token: %w", err)
}
return &IBDCookie{
Token: string(token),
Expiry: expiry,
}, nil
}
func GetCookies(
ctx context.Context,
exec Executor,
kms keys.KeyManagementService,
subject string,
degraded bool,
) ([]IBDCookie, error) {
rows, err := exec.QueryContext(ctx, `
SELECT ibd_tokens.id, token, encrypted_key, kms_key_name, expires_at
FROM ibd_tokens
INNER JOIN keys ON encryption_key = keys.id
WHERE user_subject = $1
AND expires_at > NOW()
AND degraded = $2
ORDER BY expires_at DESC;`, subject, degraded)
if err != nil {
return nil, fmt.Errorf("unable to get ibd cookies: %w", err)
}
cookies := make([]IBDCookie, 0)
for rows.Next() {
var id uint
var encryptedToken, encryptedKey []byte
var keyName string
var expiry time.Time
err = rows.Scan(&id, &encryptedToken, &encryptedKey, &keyName, &expiry)
if err != nil {
return nil, fmt.Errorf("unable to scan sql row into ibd cookie: %w", err)
}
// Set the expiry to UTC explicitly.
// For some reason, the expiry time is set to location="".
expiry = expiry.UTC()
token, err := keys.Decrypt(ctx, kms, keyName, encryptedToken, encryptedKey)
if err != nil {
return nil, fmt.Errorf("unable to decrypt token: %w", err)
}
cookie := IBDCookie{
ID: id,
Token: string(token),
Expiry: expiry,
}
cookies = append(cookies, cookie)
}
return cookies, nil
}
func AddCookie(
ctx context.Context,
exec TransactionExecutor,
kms keys.KeyManagementService,
subject string,
cookie *http.Cookie,
) error {
tx, err := exec.BeginTx(ctx, nil)
if err != nil {
return err
}
// Get the key ID for the user
user, err := GetUser(ctx, tx, subject)
if err != nil {
return fmt.Errorf("unable to get user: %w", err)
}
if user.EncryptionKeyID == nil {
return errors.New("user does not have an encryption key")
}
// Get the key
var keyName string
var key []byte
err = tx.QueryRowContext(ctx, `
SELECT kms_key_name, encrypted_key
FROM keys
WHERE id = $1;`,
*user.EncryptionKeyID,
).Scan(&keyName, &key)
if err != nil {
return fmt.Errorf("unable to get key: %w", err)
}
// Encrypt the token
encryptedToken, err := keys.EncryptWithKey(ctx, kms, keyName, key, []byte(cookie.Value))
if err != nil {
return fmt.Errorf("unable to encrypt token: %w", err)
}
// Add the cookie to the database
_, err = exec.ExecContext(ctx, `
INSERT INTO ibd_tokens (token, expires_at, user_subject, encryption_key)
VALUES ($1, $2, $3, $4)`, encryptedToken, cookie.Expires, subject, *user.EncryptionKeyID)
if err != nil {
return fmt.Errorf("unable to add cookie: %w", err)
}
return nil
}
func ReportCookieFailure(ctx context.Context, exec Executor, id uint) error {
_, err := exec.ExecContext(ctx, `
UPDATE ibd_tokens
SET degraded = TRUE
WHERE id = $1;`, id)
if err != nil {
return fmt.Errorf("unable to report cookie failure: %w", err)
}
return nil
}
func RepairCookie(ctx context.Context, exec Executor, id uint) error {
_, err := exec.ExecContext(ctx, `
UPDATE ibd_tokens
SET degraded = FALSE
WHERE id = $1;`, id)
if err != nil {
return fmt.Errorf("unable to report cookie failure: %w", err)
}
return nil
}
type IBDCookie struct {
ID uint
Token string
Expiry time.Time
}
func (c *IBDCookie) ToHTTPCookie() *http.Cookie {
return &http.Cookie{
Name: ".ASPXAUTH",
Value: c.Token,
Path: "/",
Domain: "investors.com",
Expires: c.Expiry,
Secure: true,
HttpOnly: false,
SameSite: http.SameSiteLaxMode,
}
}
|
