@@ -14,12 +14,12 @@ import (
1414 "github.com/rs/zerolog/log"
1515 "github.com/samber/do"
1616 "github.com/sethvargo/go-diceware/diceware"
17+ "github.com/therealpaulgg/ssh-sync-common/pkg/dto"
18+ "github.com/therealpaulgg/ssh-sync-common/pkg/wsutils"
1719 pqc "github.com/therealpaulgg/ssh-sync-server/pkg/crypto"
1820 "github.com/therealpaulgg/ssh-sync-server/pkg/database/models"
1921 "github.com/therealpaulgg/ssh-sync-server/pkg/database/repository"
2022 "github.com/therealpaulgg/ssh-sync-server/pkg/web/middleware/context_keys"
21- "github.com/therealpaulgg/ssh-sync/pkg/dto"
22- "github.com/therealpaulgg/ssh-sync/pkg/utils"
2323)
2424
2525// Computer A creates a live connection.
@@ -38,7 +38,7 @@ type ChallengeResponse struct {
3838type ChallengeSession struct {
3939 Username string
4040 ChallengeAccepted chan bool
41- ChallengerChannel chan [] byte
41+ ChallengerChannel chan * dto. PublicKeyDto
4242 ResponderChannel chan []byte
4343}
4444
@@ -85,15 +85,15 @@ func MachineChallengeResponseHandler(i *do.Injector, r *http.Request, w http.Res
8585 log .Warn ().Msg ("Could not get user from context" )
8686 return
8787 }
88- foo , err := utils .ReadClientMessage [dto.ChallengeResponseDto ](& conn )
88+ foo , err := wsutils .ReadClientMessage [dto.ChallengeResponseDto ](& conn )
8989 if err != nil {
9090 log .Err (err ).Msg ("Error reading client message" )
9191 return
9292 }
9393 chalChan , ok := ChallengeResponseDict .ReadChallenge (foo .Data .Challenge )
9494 if ! ok {
9595 log .Warn ().Msg ("Could not find challenge in dict" )
96- if err := utils .WriteServerError [dto.ChallengeSuccessEncryptedKeyDto ](& conn , "Invalid challenge response." ); err != nil {
96+ if err := wsutils .WriteServerError [dto.ChallengeSuccessEncryptedKeyDto ](& conn , "Invalid challenge response." ); err != nil {
9797 log .Err (err ).Msg ("Error writing server error" )
9898 }
9999 return
@@ -109,19 +109,20 @@ func MachineChallengeResponseHandler(i *do.Injector, r *http.Request, w http.Res
109109 key := <- chalChan .ChallengerChannel
110110 if key == nil {
111111 log .Debug ().Msg ("Response from challenger channel - key is nil. Exiting." )
112- if err := utils .WriteServerError [dto.ChallengeSuccessEncryptedKeyDto ](& conn , "Error responding to challenge - client abruptly closed connection." ); err != nil {
112+ if err := wsutils .WriteServerError [dto.ChallengeSuccessEncryptedKeyDto ](& conn , "Error responding to challenge - client abruptly closed connection." ); err != nil {
113113 log .Err (err ).Msg ("Error writing server error" )
114114 }
115115 return
116116 }
117117 keys := dto.ChallengeSuccessEncryptedKeyDto {
118- PublicKey : key ,
118+ PublicKey : key .PublicKey ,
119+ EncapsulationKey : key .EncapsulationKey ,
119120 }
120- if err := utils .WriteServerMessage (& conn , keys ); err != nil {
121+ if err := wsutils .WriteServerMessage (& conn , keys ); err != nil {
121122 log .Err (err ).Msg ("Error writing server message" )
122123 return
123124 }
124- encMasterKeyDto , err := utils .ReadClientMessage [dto.EncryptedMasterKeyDto ](& conn )
125+ encMasterKeyDto , err := wsutils .ReadClientMessage [dto.EncryptedMasterKeyDto ](& conn )
125126 if err != nil {
126127 log .Err (err ).Msg ("Error reading client message" )
127128 return
@@ -140,17 +141,17 @@ func NewMachineChallenge(i *do.Injector, r *http.Request, w http.ResponseWriter)
140141
141142func NewMachineChallengeHandler (i * do.Injector , r * http.Request , w http.ResponseWriter , c * net.Conn ) {
142143 conn := * c
143- defer conn .Close ()
144+ defer conn .Close ()
144145 // first message sent should be JSON payload
145- userMachine , err := utils .ReadClientMessage [dto.UserMachineDto ](& conn )
146+ userMachine , err := wsutils .ReadClientMessage [dto.UserMachineDto ](& conn )
146147 if err != nil {
147148 log .Err (err ).Msg ("Error reading client message" )
148149 return
149150 }
150151 userRepo := do.MustInvoke [repository.UserRepository ](i )
151152 user , err := userRepo .GetUserByUsername (userMachine .Data .Username )
152153 if errors .Is (err , sql .ErrNoRows ) || user == nil {
153- if err := utils .WriteServerError [dto.MessageDto ](& conn , "User not found" ); err != nil {
154+ if err := wsutils .WriteServerError [dto.MessageDto ](& conn , "User not found" ); err != nil {
154155 log .Err (err ).Msg ("Error writing server error" )
155156 }
156157 return
@@ -163,7 +164,7 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
163164 machine , err := machineRepo .GetMachineByNameAndUser (userMachine .Data .MachineName , user .ID )
164165 // if the machine already exists, reject
165166 if err == nil && machine .ID != uuid .Nil {
166- if err = utils .WriteServerError [dto.MessageDto ](& conn , "Machine already exists" ); err != nil {
167+ if err = wsutils .WriteServerError [dto.MessageDto ](& conn , "Machine already exists" ); err != nil {
167168 log .Err (err ).Msg ("Error writing server error" )
168169 }
169170 return
@@ -179,13 +180,13 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
179180 words , err := diceware .GenerateWithWordList (3 , diceware .WordListEffLarge ())
180181 if err != nil {
181182 log .Err (err ).Msg ("Error generating diceware" )
182- if err := utils .WriteServerError [dto.MessageDto ](& conn , "Error generating diceware" ); err != nil {
183+ if err := wsutils .WriteServerError [dto.MessageDto ](& conn , "Error generating diceware" ); err != nil {
183184 log .Err (err ).Msg ("Error writing server error" )
184185 }
185186 return
186187 }
187188 challengePhrase := strings .Join (words , "-" )
188- if err := utils .WriteServerMessage (& conn , dto.MessageDto {Message : challengePhrase }); err != nil {
189+ if err := wsutils .WriteServerMessage (& conn , dto.MessageDto {Message : challengePhrase }); err != nil {
189190 log .Err (err ).Msg ("Error writing challenge phrase" )
190191 return
191192 }
@@ -199,7 +200,7 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
199200 ChallengeResponseDict .WriteChallenge (challengePhrase , ChallengeSession {
200201 Username : user .Username ,
201202 ChallengeAccepted : make (chan bool ),
202- ChallengerChannel : make (chan [] byte ),
203+ ChallengerChannel : make (chan * dto. PublicKeyDto ),
203204 ResponderChannel : make (chan []byte ),
204205 })
205206 defer func () {
@@ -253,16 +254,16 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
253254 challengeResult := <- challengeResponse
254255
255256 if ! challengeResult {
256- if err := utils .WriteServerError [dto.MessageDto ](& conn , "Challenge timed out" ); err != nil {
257+ if err := wsutils .WriteServerError [dto.MessageDto ](& conn , "Challenge timed out" ); err != nil {
257258 log .Err (err ).Msg ("Error writing server error" )
258259 }
259260 return
260261 }
261- if err := utils .WriteServerMessage (& conn , dto.MessageDto {Message : "Challenge accepted!" }); err != nil {
262+ if err := wsutils .WriteServerMessage (& conn , dto.MessageDto {Message : "Challenge accepted!" }); err != nil {
262263 log .Err (err ).Msg ("Error writing challenge accepted" )
263264 return
264265 }
265- pubkey , err := utils .ReadClientMessage [dto.PublicKeyDto ](& conn )
266+ pubkey , err := wsutils .ReadClientMessage [dto.PublicKeyDto ](& conn )
266267 if err != nil {
267268 log .Err (err ).Msg ("Error reading client message" )
268269 return
@@ -271,23 +272,23 @@ func NewMachineChallengeHandler(i *do.Injector, r *http.Request, w http.Response
271272 // Validate the public key format before storing
272273 if _ , err := pqc .ValidatePublicKey (pubkey .Data .PublicKey ); err != nil {
273274 log .Err (err ).Msg ("Invalid public key format in challenge flow" )
274- if err := utils .WriteServerError [dto.MessageDto ](& conn , "Invalid public key format" ); err != nil {
275+ if err := wsutils .WriteServerError [dto.MessageDto ](& conn , "Invalid public key format" ); err != nil {
275276 log .Err (err ).Msg ("Error writing server error" )
276277 }
277278 return
278279 }
279- cha .ChallengerChannel <- pubkey .Data . PublicKey
280+ cha .ChallengerChannel <- & pubkey .Data
280281 encryptedMasterKey := <- cha .ResponderChannel
281282 machine .PublicKey = pubkey .Data .PublicKey
282283 if _ , err = machineRepo .CreateMachine (machine ); err != nil {
283284 log .Err (err ).Msg ("Error creating machine" )
284285 return
285286 }
286- if err := utils .WriteServerMessage (& conn , dto.EncryptedMasterKeyDto {EncryptedMasterKey : encryptedMasterKey }); err != nil {
287+ if err := wsutils .WriteServerMessage (& conn , dto.EncryptedMasterKeyDto {EncryptedMasterKey : encryptedMasterKey }); err != nil {
287288 log .Err (err ).Msg ("Error writing encrypted master key" )
288289 return
289290 }
290- if err := utils .WriteServerMessage (& conn , dto.MessageDto {Message : "Everything is done, you can now use ssh-sync" }); err != nil {
291+ if err := wsutils .WriteServerMessage (& conn , dto.MessageDto {Message : "Everything is done, you can now use ssh-sync" }); err != nil {
291292 log .Err (err ).Msg ("Error writing final message" )
292293 return
293294 }
0 commit comments