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
|
package database
import (
"context"
"database/sql"
"errors"
"fmt"
"net/http"
"time"
"ibd-trader/internal/keys"
)
type CookieStore interface {
CookieSource
AddCookie(ctx context.Context, subject string, cookie *http.Cookie) error
RepairCookie(ctx context.Context, id uint) error
}
type CookieSource interface {
GetAnyCookie(ctx context.Context) (*IBDCookie, error)
GetCookies(ctx context.Context, subject string, degraded bool) ([]IBDCookie, error)
ReportCookieFailure(ctx context.Context, id uint) error
}
func (d *database) GetAnyCookie(ctx context.Context) (*IBDCookie, error) {
row, err := d.queryRow(ctx, d.db, "cookies/get_any_cookie")
if err != nil {
return nil, fmt.Errorf("unable to get any ibd cookie: %w", err)
}
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)
}
token, err := keys.Decrypt(ctx, d.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 (d *database) GetCookies(ctx context.Context, subject string, degraded bool) ([]IBDCookie, error) {
row, err := d.query(ctx, d.db, "cookies/get_cookies", subject, degraded)
if err != nil {
return nil, fmt.Errorf("unable to get ibd cookies: %w", err)
}
cookies := make([]IBDCookie, 0)
for row.Next() {
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 {
return nil, fmt.Errorf("unable to scan sql row into ibd cookie: %w", err)
}
token, err := keys.Decrypt(ctx, d.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 (d *database) AddCookie(ctx context.Context, subject string, cookie *http.Cookie) error {
// Get the key ID for the user
user, err := d.GetUser(ctx, 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
key, err := d.GetKey(ctx, *user.EncryptionKeyID)
if err != nil {
return fmt.Errorf("unable to get key: %w", err)
}
// Encrypt the token
encryptedToken, err := keys.EncryptWithKey(ctx, d.kms, key.Name, key.Key, []byte(cookie.Value))
if err != nil {
return fmt.Errorf("unable to encrypt token: %w", err)
}
// Add the cookie to the database
_, err = d.exec(ctx, d.db, "cookies/add_cookie", encryptedToken, cookie.Expires, subject, key.Id)
if err != nil {
return fmt.Errorf("unable to add cookie: %w", err)
}
return nil
}
func (d *database) ReportCookieFailure(ctx context.Context, id uint) error {
_, err := d.exec(ctx, d.db, "cookies/set_cookie_degraded", true, id)
if err != nil {
return fmt.Errorf("unable to report cookie failure: %w", err)
}
return nil
}
func (d *database) RepairCookie(ctx context.Context, id uint) error {
_, err := d.exec(ctx, d.db, "cookies/set_cookie_degraded", false, 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,
}
}
|
