6
6
"net"
7
7
"net/http"
8
8
"strings"
9
+ "sync"
9
10
"time"
10
11
11
12
"github.com/gobwas/ws"
@@ -40,8 +41,30 @@ type Something struct {
40
41
ResponderChannel chan []byte
41
42
}
42
43
44
+ type SafeChallengeResponseDict struct {
45
+ mux sync.Mutex
46
+ dict map [string ]Something
47
+ }
48
+
49
+ // Utility method for safely writing to the dict
50
+ func (c * SafeChallengeResponseDict ) WriteChallenge (challengePhrase string , data Something ) {
51
+ c .mux .Lock ()
52
+ c .dict [challengePhrase ] = data
53
+ c .mux .Unlock ()
54
+ }
55
+
56
+ // Utility method for safely reading from the dict
57
+ func (c * SafeChallengeResponseDict ) ReadChallenge (challengePhrase string ) (Something , bool ) {
58
+ c .mux .Lock ()
59
+ data , exists := c .dict [challengePhrase ]
60
+ c .mux .Unlock ()
61
+ return data , exists
62
+ }
63
+
43
64
var ChallengeResponseChannel = make (chan ChallengeResponse )
44
- var ChallengeResponseDict = make (map [string ]Something )
65
+ var ChallengeResponseDict = SafeChallengeResponseDict {
66
+ dict : make (map [string ]Something ),
67
+ }
45
68
46
69
func MachineChallengeResponse (i * do.Injector , r * http.Request , w http.ResponseWriter ) error {
47
70
conn , _ , _ , err := ws .UpgradeHTTP (r , w )
@@ -65,7 +88,7 @@ func MachineChallengeResponseHandler(i *do.Injector, r *http.Request, w http.Res
65
88
log .Err (err ).Msg ("Error reading client message" )
66
89
return
67
90
}
68
- chalChan , ok := ChallengeResponseDict [ foo .Data .Challenge ]
91
+ chalChan , ok := ChallengeResponseDict . ReadChallenge ( foo .Data .Challenge )
69
92
if ! ok {
70
93
log .Warn ().Msg ("Could not find challenge in dict" )
71
94
if err := utils .WriteServerError [dto.ChallengeSuccessEncryptedKeyDto ](& conn , "Invalid challenge response." ); err != nil {
@@ -164,34 +187,63 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
164
187
// Computer A will receive the public key, decrypt the master key, encrypt the master key with the public key, and send it back to the server.
165
188
// At this point Computer B will be able to communicate freely.
166
189
167
- ChallengeResponseDict [ challengePhrase ] = Something {
190
+ ChallengeResponseDict . WriteChallenge ( challengePhrase , Something {
168
191
Username : user .Username ,
169
192
ChallengeAccepted : make (chan bool ),
170
193
ChallengerChannel : make (chan []byte ),
171
194
ResponderChannel : make (chan []byte ),
172
- }
195
+ })
173
196
defer func () {
174
- close (ChallengeResponseDict [challengePhrase ].ChallengeAccepted )
175
- close (ChallengeResponseDict [challengePhrase ].ChallengerChannel )
176
- close (ChallengeResponseDict [challengePhrase ].ResponderChannel )
177
- delete (ChallengeResponseDict , challengePhrase )
197
+ ChallengeResponseDict .mux .Lock ()
198
+ defer ChallengeResponseDict .mux .Unlock ()
199
+ item , exists := ChallengeResponseDict .dict [challengePhrase ]
200
+ if exists {
201
+ close (item .ChallengeAccepted )
202
+ close (item .ChallengerChannel )
203
+ close (item .ResponderChannel )
204
+ delete (ChallengeResponseDict .dict , challengePhrase )
205
+ }
178
206
}()
179
207
timer := time .NewTimer (30 * time .Second )
180
208
go func () {
209
+ var challengeAcceptedChan chan bool
210
+
211
+ // Lock before accessing ChallengeResponseDict
212
+ ChallengeResponseDict .mux .Lock ()
213
+ if item , exists := ChallengeResponseDict .dict [challengePhrase ]; exists {
214
+ challengeAcceptedChan = item .ChallengeAccepted
215
+ }
216
+ ChallengeResponseDict .mux .Unlock ()
217
+
218
+ // If the channel does not exist, return to avoid a nil channel operation
219
+ if challengeAcceptedChan == nil {
220
+ return
221
+ }
222
+
181
223
for {
182
224
select {
183
225
case <- timer .C :
184
- ChallengeResponseDict [challengePhrase ].ChallengeAccepted <- false
226
+ ChallengeResponseDict .mux .Lock ()
227
+ // Check if the challenge still exists before sending to the channel
228
+ if _ , exists := ChallengeResponseDict .dict [challengePhrase ]; exists {
229
+ ChallengeResponseDict .dict [challengePhrase ].ChallengeAccepted <- false
230
+ }
231
+ ChallengeResponseDict .mux .Unlock ()
185
232
return
186
- case chalWon := <- ChallengeResponseDict [ challengePhrase ]. ChallengeAccepted :
233
+ case chalWon := <- challengeAcceptedChan :
187
234
if chalWon {
188
235
timer .Stop ()
189
236
}
190
237
return
191
238
}
192
239
}
193
240
}()
194
- challengeResult := <- ChallengeResponseDict [challengePhrase ].ChallengeAccepted
241
+ cha , ok := ChallengeResponseDict .ReadChallenge (challengePhrase )
242
+ if ! ok {
243
+ log .Err (err ).Msg ("Error getting challenge from dict" )
244
+ return
245
+ }
246
+ challengeResult := <- cha .ChallengeAccepted
195
247
196
248
if ! challengeResult {
197
249
if err := utils .WriteServerError [dto.MessageDto ](& conn , "Challenge timed out" ); err != nil {
@@ -208,8 +260,9 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
208
260
log .Err (err ).Msg ("Error reading client message" )
209
261
return
210
262
}
211
- ChallengeResponseDict [challengePhrase ].ChallengerChannel <- pubkey .Data .PublicKey
212
- encryptedMasterKey := <- ChallengeResponseDict [challengePhrase ].ResponderChannel
263
+
264
+ cha .ChallengerChannel <- pubkey .Data .PublicKey
265
+ encryptedMasterKey := <- cha .ResponderChannel
213
266
machine .PublicKey = pubkey .Data .PublicKey
214
267
if _ , err = machineRepo .CreateMachine (machine ); err != nil {
215
268
log .Err (err ).Msg ("Error creating machine" )
0 commit comments