summaryrefslogblamecommitdiffstats
path: root/src/wireguard.c
blob: 53e2cdb6ceed928433f32005f52ef93e0f395d6f (plain) (tree)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16















                                                                           
                      
                       
                      
                      
                       
                     
                           
                                
 


                              
 
                          
 
                                   



                                                               

                                                   



                                    
 


                                                                             



                                                  
 





                                                               
 
                                                              
                                                                
 















                                                                                           


                                                              
                                                         
                                          
 
                     
    
                                              
                                        

                                                                  
 
                                 

                              
                                  
                                    

                                           
                                   

                                                        
 


                                                               
                               
                                                  



                               
                                                              
 


                               




                                        

                              
                                 


                                              



                                     




                                                                           
 


                              
                                                       
 
                                
                                                    
                                
 
                                                                
                    

 
                   
                                         
 
                                                 


                                                                        








                                                           
 
                                   
                                                                              
                                   
 


                                             



                       
                                                                   
 

                                                                        
                                    
                              

                                                     


                                                                          


                                    
         
                 
                                                            
 
                    


    
                                 
 
                                                     

 
    
                                  
 
                                                      
                                        
                            


                                            
    
                                          
 
                                                                            

 
    
                                           
 
                                                                             
                                                  


    

                                            
                                
                             
                                


    


                                            
 
                                






                                        
                                










                                        
                                              
 
                                       



                                                        
                              
                                                                             
                                   
                                                         




                                                                              
                
                                           
                       
         

                                   

                                   
 
                                
                                             



                                                            
                                


                                             

                                       

 

                                                               
 
                                
                              
                                


    
                                                               
 
                                
                              
                                

 

                                            
 
                                       

                                       
                                                    
                                
                   

 


                                          
                        
                                   
 

                                       
                                                    
                                




                                                                                
                                                                     



                                

 



                                                                              
 
                             
                             


                                
                                   
 
                                  
 

                  


                                                                        
                                   
 
                             

                                                            
 


                                                                             
 
                                                

                                                                         
                                                 
 


                                                                  
 
                                                                     
                                            
 


                                                             
 
                                                                   
                                                                          
 

                                                                      
 
                               

                                                               
 

                                                                          

                                                                 
 






                                                                             
 



                                                                   
 


                                            
                                            


                                                    
 
                                                                
                          
                     
      

                                        
                                        


                                        
 
             

                                                                          
 
                             
                               
                                 

                                   
                                  
 
                                                                            
                             
 
                                  
                                   
 
                                   
                                  




                                           
 
                                  
                                                            
 

                                                                         
 
                                                         
                                                 
 
                                  
 

                                                                                
 
                                                            
 
                                                            


                                                                    
 
                          

                                                               
 









                                                            
 
                                    
 
                     
      

                                        

                                                
                                        
                   
 

 
             

                                                                        
 


                                      
 
                                  
 
                                                                              
                             
 

                                                      
 



                                                                           
                                      
 

                                              


                                                       
                                           
 
                     
      

                                        
                   

 


                                                                           
 
                                   


                                                                        
 
                                                                           
                             
 

                                    
                                                                    


                                                     
 


                                                                              
 


                                                                              
 

                                                                
 
                                                
                                                              
 
                                                     



                                                                     
                                                                                       
 

                     

                                        
                   
 
 
 
             

                                                                              
 
                             


                                          
 
                                  
 
                                                          
                             
 
                                
                                                                              
                                 
                                                       








                                                      
 


                                                                               
                                            
 
                                   
                                         
 
                             

                                           
                                                               
 


                                                                             
 



                                                                            
 
                                                             

                         

                                                                           

                                          
 


                                                                                       
 


                                                                      
 
                                        






                                                                                 



                                                    
 
                     
      


                                           

                                        
                   

 
             

                                                                          
 
                               

                                   

                                  
                                                                
                             
 


                                              
 


                                   
 


                                                                            
 
                                                               
 

                                                                       
 
                                                     
 
                                                        
 
                                                            
 



                                                                                             
 
                                        
                                                                                        

                                               
                                                                                   
                                                                               
                                                              
 


                                                          

                                   



                                                          
                     
      

                                        
                   

 
             

                                                                         
 












                                                                          
 
 



                                                                           
                                   
 

                                  
                                                                
                             
 
                                            


                                                     
 

                                                                              

                                     

                                             
                                                                
 

                                                                                  

                                               
                                                       
                                                                              


                                                                       
                                                                                       
 
                                           
                     
      

                                        
                   

 




                                                                           
 

                                                 
 
                                                                               
 

                       
 


                                                                           
 

                       
 


                                                                                          
 

                       
 



                                                                                          
 







                                                                                                 
 
 
    
                                                                  
 

                                                                   
 







                                                                                        
 













                                                                                                
 



                                                                                                

 
    
                                               
 



                                                                        

 










                                                             
 
                
                                     
 
                                   
 




                                                   
                                                                           
                                         
                                                                              
                                       
                                                                            
                                     
                                                                               
                                        
            
                                      
 
 

















                                                                        

    






                                                       
                                                                                

                                 
                                                            
                                                                      
                                                                             


             
                                                                            

                                 
                                                            
                                                                    
                                                                             
 
 










                                                                              
 

                     
















                                                                            










                                   





                                                              





                                                      
    





                                                                         








































                                                                                       
/*
 * Copyright (c) 2019 Matt Dunwoodie <ncon@noconroy.net>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <sys/types.h>
#include <sys/malloc.h>
#include <sys/systm.h>
#include <sys/mutex.h>
#include <sys/refcnt.h>
#include <sys/time.h>
#include <sys/antireplay.h>
#include <lib/libkern/libkern.h>

#include <crypto/blake2s.h>
#include <crypto/curve25519.h>
#include <crypto/chachapoly.h>

#include <net/wireguard.h>

/* Constants for the wg protocol */
#define WG_CONSTRUCTION "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
#define WG_IDENTIFIER "WireGuard v1 zx2c4 Jason@zx2c4.com"
#define WG_MAC1 "mac1----"
#define WG_COOKIE "cookie--"

/* First byte indicating packet type on the wire */
#define WG_MSG_INITIATION htole32(1)
#define WG_MSG_RESPONSE htole32(2)
#define WG_MSG_COOKIE htole32(3)
#define WG_MSG_TRANSPORT htole32(4)

/* Ensure struct sizes are correct at compile time. This allows us to overlay
 * the struct on the incoming network packet and access it as if it was a
 * struct. */
CTASSERT(sizeof(struct wg_msg_initiation) == 148);
CTASSERT(sizeof(struct wg_msg_response) == 92);
CTASSERT(sizeof(struct wg_msg_cookie) == 64);
CTASSERT(sizeof(struct wg_msg_transport) == 16);

CTASSERT(WG_KEY_SIZE == CHACHA20POLY1305_KEY_SIZE);
CTASSERT(WG_HASH_SIZE == BLAKE2S_HASH_SIZE);
CTASSERT(WG_KEY_SIZE == BLAKE2S_KEY_SIZE);
CTASSERT(WG_KEY_SIZE == WG_HASH_SIZE);
CTASSERT(WG_MSG_PADDING_SIZE == CHACHA20POLY1305_AUTHTAG_SIZE);
CTASSERT(WG_XNONCE_SIZE == XCHACHA20POLY1305_NONCE_SIZE);

#define ret_error(err) do { ret = err; goto leave; } while (0)
#define offsetof(type, member)  ((size_t)(&((type *)0)->member))

/* wireguard.c */
void wg_kdf(uint8_t [WG_HASH_SIZE], uint8_t [WG_HASH_SIZE],
    uint8_t [WG_HASH_SIZE], uint8_t [WG_KEY_SIZE], uint8_t *, size_t);
void wg_hash2(uint8_t [WG_HASH_SIZE], uint8_t *, size_t, uint8_t *, size_t);
void wg_mix_hash(struct wg_handshake *, uint8_t *, size_t);
void wg_mix_dh(struct wg_handshake *, uint8_t [WG_KEY_SIZE], uint8_t [WG_KEY_SIZE]);
void wg_mix_psk(struct wg_handshake *, uint8_t [WG_KEY_SIZE]);
void wg_handshake_encrypt(struct wg_handshake *, uint8_t *, uint8_t *, size_t);
int wg_handshake_decrypt(struct wg_handshake *, uint8_t *, uint8_t *, size_t);
void wg_timestamp_get(uint8_t [WG_TIMESTAMP_SIZE]);
int wg_timespec_timedout(struct timespec *, time_t);
enum wg_pkt_type wg_pkt_type(uint8_t *, size_t);
void wg_keypair_generate(struct wg_keypair *);
enum wg_error wg_msg_initiation_valid_mac2(struct wg_msg_initiation *, struct wg_cookie *);
enum wg_error wg_msg_response_valid_mac2(struct wg_msg_response *, struct wg_cookie *);

struct wg_session	*wg_peer_hs_session(struct wg_peer *);
struct wg_session	*wg_peer_ks_session(struct wg_peer *);

struct wg_session *wg_peer_new_session(struct wg_peer *);
void wg_session_drop(struct wg_session *);

/* Some crappy API */
void
wg_device_init(struct wg_device *dev, int ipl,
    void (*notify_fn)(struct wg_peer *),
    void (*outq_fn)(struct wg_peer *, enum wg_pkt_type, uint32_t),
    void (*cleanup_fn)(struct wg_peer *), void *arg)
{
	bzero(dev, sizeof(*dev));
	dev->d_arg = arg;
	dev->d_outq = outq_fn;
	dev->d_notify = notify_fn;
	dev->d_cleanup = cleanup_fn;
	fm_init(&dev->d_peers, 4, ipl);
	fm_init(&dev->d_sessions, 12, ipl);
	mtx_init(&dev->d_mtx, ipl);
	/* d_cookie_maker, d_keypair initialised to 0 */
}

void
wg_device_setkey(struct wg_device *dev, struct wg_privkey *key)
{
	mtx_enter(&dev->d_mtx);
	wg_keypair_from_key(&dev->d_keypair, key);
	mtx_leave(&dev->d_mtx);
}

void
wg_device_getkey(struct wg_device *dev, struct wg_keypair *kp)
{
	mtx_enter(&dev->d_mtx);
	*kp = dev->d_keypair;
	mtx_leave(&dev->d_mtx);
}

void
wg_device_destroy(struct wg_device *dev)
{
	struct map_item *item;

	/* TODO lock fixed map */
	FM_FOREACH_FILLED(item, &dev->d_peers)
		wg_peer_drop(item->value);

	fm_destroy(&dev->d_sessions);
	fm_destroy(&dev->d_peers);
}

struct wg_peer *
wg_device_new_peer(struct wg_device *dev, struct wg_pubkey *key, void *arg)
{
	struct wg_peer *peer;
	peer = malloc(sizeof(*peer), M_DEVBUF, M_WAITOK | M_ZERO);

	peer->p_arg = arg;
	peer->p_device = dev;
	peer->p_remote = *key;
	mtx_init(&peer->p_mtx, dev->d_mtx.mtx_wantipl);

	mtx_enter(&peer->p_mtx);
	peer->p_id = fm_insert(&dev->d_peers, peer);
	mtx_leave(&peer->p_mtx);

	/* All other elements of wg_peer are nulled by M_ZERO */
	return peer;
}

struct wg_session *
wg_peer_new_session(struct wg_peer *peer)
{
	struct wg_session *session, *old_session;
	session = malloc(sizeof(*session), M_DEVBUF, M_WAITOK | M_ZERO);

	getnanotime(&session->s_created);
	session->s_peer = peer;
	session->s_state = WG_STATE_NEW;
	mtx_init(&session->s_mtx, peer->p_mtx.mtx_wantipl);

	mtx_enter(&peer->p_mtx);

	old_session = peer->p_hs_session;
	peer->p_hs_session = session;
	mtx_leave(&peer->p_mtx);

	mtx_enter(&session->s_mtx);
	session->s_local_id = fm_insert(&peer->p_device->d_sessions, session);
	mtx_leave(&session->s_mtx);

	if (old_session)
		wg_session_drop(old_session);

	return session;
}

struct wg_peer *
wg_device_ref_peerkey(struct wg_device *dev, struct wg_pubkey *key)
{
	/* For the time being, we just iterate through peers to find the
	 * matching peer */
	struct wg_peer *peer = NULL;
	struct map_item *item;

	/* TODO lock, or use better data structure */
	FM_FOREACH_FILLED(item, &dev->d_peers) {
		peer = item->value;
		if (memcmp(key->k, peer->p_remote.k, sizeof(key->k)) == 0)
			break;
		else
			peer = NULL;
	}
	if (peer)
		peer = fm_lookup(&dev->d_peers, peer->p_id);

	return peer;
}

void
wg_peer_put(struct wg_peer *peer)
{
	fm_put(&peer->p_device->d_peers, peer->p_id);
}

void
wg_peer_drop(struct wg_peer *peer)
{
	fm_drop(&peer->p_device->d_peers, peer->p_id);
	peer->p_device->d_cleanup(peer);
	wg_peer_clean(peer);
	free(peer, M_DEVBUF, sizeof(*peer));
}

void
wg_session_put(struct wg_session *session)
{
	fm_put(&session->s_peer->p_device->d_sessions, session->s_local_id);
}

void
wg_session_drop(struct wg_session *session)
{
	fm_drop(&session->s_peer->p_device->d_sessions, session->s_local_id);
	free(session, M_DEVBUF, sizeof(*session));
}

void
wg_peer_reset_attempts(struct wg_peer *peer)
{
	mtx_enter(&peer->p_mtx);
	peer->p_attempts = 0;
	mtx_leave(&peer->p_mtx);
}

void
wg_peer_clean(struct wg_peer *peer)
{
	struct wg_session *hs, *ks, *ks_old;

	mtx_enter(&peer->p_mtx);
	hs = peer->p_hs_session;
	ks = peer->p_ks_session;
	ks_old = peer->p_ks_session_old;

	peer->p_hs_session = NULL;
	peer->p_ks_session = NULL;
	peer->p_ks_session_old = NULL;
	mtx_leave(&peer->p_mtx);

	if (hs != NULL)
		wg_session_drop(hs);
	if (ks != NULL)
		wg_session_drop(ks);
	if (ks_old != NULL)
		wg_session_drop(ks_old);

}

void
wg_session_promote(struct wg_session *session)
{
	struct wg_session *old_session;
	struct wg_peer *peer = session->s_peer;
	struct wg_keyset *ks = &session->s_keyset;
	struct wg_handshake *hs = &session->s_handshake;

	/* TODO make better */
	/* Setup session: derive keys, initialise the antireplay structure */
	mtx_enter(&session->s_mtx);
	if (session->s_state == WG_STATE_RECV_RESPONSE) {
		session->s_state = WG_STATE_INITIATOR;
		wg_kdf(ks->k_txkey.k, ks->k_rxkey.k, NULL, hs->h_ck, NULL, 0);
	} else if (session->s_state == WG_STATE_MADE_RESPONSE) {
		session->s_state = WG_STATE_RESPONDER;
		wg_kdf(ks->k_rxkey.k, ks->k_txkey.k, NULL, hs->h_ck, NULL, 0);
	} else {
		mtx_leave(&session->s_mtx);
		return;
	}

	antireplay_init(&ks->k_ar);
	mtx_leave(&session->s_mtx);


	mtx_enter(&peer->p_mtx);
	old_session = peer->p_ks_session_old;
	if (peer->p_ks_session != NULL)
		peer->p_ks_session_old = peer->p_ks_session;
	peer->p_ks_session = peer->p_hs_session;
	peer->p_hs_session = NULL;
	mtx_leave(&peer->p_mtx);

	if (old_session != NULL)
		wg_session_drop(old_session);

	peer->p_device->d_notify(peer);
}

void
wg_peer_setshared(struct wg_peer *peer, struct wg_privkey *key)
{
	mtx_enter(&peer->p_mtx);
	peer->p_shared = *key;
	mtx_leave(&peer->p_mtx);
}

void
wg_peer_getshared(struct wg_peer *peer, struct wg_privkey *key)
{
	mtx_enter(&peer->p_mtx);
	*key = peer->p_shared;
	mtx_leave(&peer->p_mtx);
}

struct timespec
wg_peer_last_handshake(struct wg_peer *peer)
{
	struct timespec ret = { 0, 0 };
	mtx_enter(&peer->p_mtx);
	if (peer->p_ks_session != NULL)
		ret = peer->p_ks_session->s_created;
	mtx_leave(&peer->p_mtx);
	return ret;
}

struct wg_session *
wg_peer_last_session(struct wg_peer *peer)
{
	uint32_t id = 0;
	struct wg_session *session;

	mtx_enter(&peer->p_mtx);
	if (peer->p_ks_session != NULL)
		id = peer->p_ks_session->s_local_id;
	mtx_leave(&peer->p_mtx);

	if ((session = fm_lookup(&peer->p_device->d_sessions, id)) == NULL)
		return NULL;

	if (!wg_timespec_timedout(&session->s_created, WG_REJECT_AFTER_TIME) && 
	    session->s_keyset.k_txcounter < WG_REJECT_AFTER_MESSAGES)
		return session;
	
	wg_session_put(session);
	return NULL;
}

/* Crypto  */
enum wg_error
wg_device_rx_initiation(struct wg_device *dev, struct wg_msg_initiation *init,
    struct wg_session **s)
{
	struct wg_peer *peer;
	struct wg_keypair kp;
	struct wg_handshake hs;
	struct wg_timestamp ts;
	struct wg_pubkey remote;
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	*s = NULL;

	/* We want to ensure that the keypair is not modified during the
	 * handshake, so we take a local copy here and bzero it before
	 * returning */
	wg_device_getkey(dev, &kp);

	/* Noise handshake */
	memcpy(hs.h_remote.k, init->ephemeral, WG_KEY_SIZE);
	wg_keypair_generate(&hs.h_local);

	wg_hash2(hs.h_ck, WG_CONSTRUCTION, strlen(WG_CONSTRUCTION), NULL, 0);
	memcpy(hs.h_hash, hs.h_ck, WG_HASH_SIZE);
	wg_mix_hash(&hs, WG_IDENTIFIER, strlen(WG_IDENTIFIER));

	wg_mix_hash(&hs, kp.pub.k, WG_KEY_SIZE);
	wg_kdf(hs.h_ck, NULL, NULL, hs.h_ck, hs.h_remote.k, WG_KEY_SIZE);
	wg_mix_hash(&hs, hs.h_remote.k, WG_KEY_SIZE);
	wg_mix_dh(&hs, kp.priv.k, hs.h_remote.k);

	if (!wg_handshake_decrypt(&hs, remote.k, init->static_pub,
	    WG_ENCRYPTED_SIZE(sizeof(remote.k))))
		ret_error(WG_DECRYPT);

	wg_mix_hash(&hs, init->static_pub, sizeof(init->static_pub));
	wg_mix_dh(&hs, kp.priv.k, remote.k);

	if (!wg_handshake_decrypt(&hs, ts.t, init->timestamp,
	    WG_ENCRYPTED_SIZE(sizeof(ts.t))))
		ret_error(WG_DECRYPT);

	wg_mix_hash(&hs, init->timestamp, sizeof(init->timestamp));
	wg_hash2(hs.h_k, WG_MAC1, strlen(WG_MAC1), kp.pub.k, WG_KEY_SIZE);

	blake2s(hs.h_mac, (void *)init, hs.h_k, sizeof(hs.h_mac),
	    offsetof(struct wg_msg_initiation, mac1), sizeof(hs.h_k));

	/* Check MAC matches */
	if (timingsafe_bcmp(hs.h_mac, init->mac1, WG_MAC_SIZE))
		ret_error(WG_MAC);

	/* Lookup peer key that was specified in the packet, as we need to
	 * know what peer this is for. */
	if ((peer = wg_device_ref_peerkey(dev, &remote)) == NULL)
		ret_error(WG_UNKNOWN_PEER);

	/* We want to ensure this packet is not replayed, so we validate that
	 * the timestamp (not necessarily representative of the real time) is
	 * greater than the last one we have received */
	mtx_enter(&peer->p_mtx);
	if (memcmp(ts.t, peer->p_timestamp.t, sizeof(ts.t)) >= 0)
		peer->p_timestamp = ts;
	mtx_leave(&peer->p_mtx);

	if (memcmp(ts.t, peer->p_timestamp.t, sizeof(ts.t)) != 0) {
		wg_peer_put(peer);
		ret_error(WG_TIMESTAMP);
	}

	session = wg_peer_new_session(peer);

	mtx_enter(&session->s_mtx);
	session->s_remote_id = init->sender;
	session->s_handshake = hs;
	session->s_state = WG_STATE_RECV_INITIATION;
	mtx_leave(&session->s_mtx);

	dev->d_outq(peer, WG_PKT_RESPONSE, session->s_local_id);
	wg_peer_put(peer);
	*s = session;
leave:
	if (ret != WG_OK)
		wg_session_put(session);
	explicit_bzero(&kp, sizeof(kp));
	explicit_bzero(&hs, sizeof(hs));
	return ret;
}

enum wg_error
wg_device_rx_response(struct wg_device *dev, struct wg_msg_response *resp,
    struct wg_session **s)
{
	struct wg_keypair kp;
	struct wg_handshake hs;
	struct wg_privkey shared;
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((session = fm_lookup(&dev->d_sessions, resp->receiver)) == NULL)
		return WG_ID;

	/* Load requried values */
	wg_device_getkey(dev, &kp);

	mtx_enter(&session->s_mtx);
	hs = session->s_handshake;
	mtx_leave(&session->s_mtx);

	mtx_enter(&session->s_peer->p_mtx);
	shared = session->s_peer->p_shared;
	mtx_leave(&session->s_peer->p_mtx);

	/* Noise recv handshake */
	memcpy(hs.h_remote.k, resp->ephemeral, WG_KEY_SIZE);

	wg_kdf(hs.h_ck, NULL, NULL, hs.h_ck, hs.h_remote.k, WG_KEY_SIZE);
	wg_mix_hash(&hs, hs.h_remote.k, WG_KEY_SIZE);

	wg_mix_dh(&hs, hs.h_local.priv.k, hs.h_remote.k);
	wg_mix_dh(&hs, kp.priv.k, hs.h_remote.k);

	wg_mix_psk(&hs, shared.k);

	if (!wg_handshake_decrypt(&hs, NULL, resp->empty, WG_ENCRYPTED_SIZE(0)))
		ret_error(WG_DECRYPT);

	wg_mix_hash(&hs, resp->empty, WG_ENCRYPTED_SIZE(0));

	wg_hash2(hs.h_k, WG_MAC1, strlen(WG_MAC1), kp.pub.k,
	    WG_KEY_SIZE);
	blake2s(hs.h_mac, (void *)resp, hs.h_k, sizeof(hs.h_mac),
	    offsetof(struct wg_msg_response, mac1), sizeof(hs.h_k));

	/* Compare macs */
	if (timingsafe_bcmp(hs.h_mac, resp->mac1, WG_MAC_SIZE))
		ret_error(WG_MAC);

	/* Update session only if we are in correct state */
	mtx_enter(&session->s_mtx);
	if (session->s_state == WG_STATE_MADE_INITIATION) {
		session->s_handshake = hs;
		session->s_remote_id = resp->sender;
		session->s_state = WG_STATE_RECV_RESPONSE;
	} else {
		ret = WG_STATE;
	}
	mtx_leave(&session->s_mtx);

	wg_session_promote(session);

	*s = session;
leave:
	if (ret != WG_OK)
		wg_session_put(session);
	explicit_bzero(&shared, sizeof(shared));
	explicit_bzero(&kp, sizeof(kp));
	explicit_bzero(&hs, sizeof(hs));
	return ret;

}

enum wg_error
wg_device_rx_cookie(struct wg_device *dev, struct wg_msg_cookie *cookie,
    struct wg_session **s)
{
	uint8_t key[WG_KEY_SIZE];
	uint8_t value[WG_COOKIE_SIZE];
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((session = fm_lookup(&dev->d_sessions, cookie->receiver)) == NULL)
		return WG_ID;

	wg_hash2(key, WG_COOKIE, strlen(WG_COOKIE),
	    session->s_peer->p_remote.k, WG_KEY_SIZE);

	/* TODO lock for h_mac? */
	if(!xchacha20poly1305_decrypt(value, cookie->value,
	    sizeof(cookie->value), session->s_handshake.h_mac, WG_MAC_SIZE,
	    cookie->nonce, key))
		ret_error(WG_DECRYPT);

	/* Update peer with new cookie data */
	mtx_enter(&session->s_peer->p_mtx);
	memcpy(session->s_peer->p_cookie.cookie, value,
	    sizeof(session->s_peer->p_cookie.cookie));
	getnanotime(&session->s_peer->p_cookie.time);
	mtx_leave(&session->s_peer->p_mtx);

	*s = session;
leave:
	if (ret != WG_OK)
		wg_session_put(session);
	return ret;
}

enum wg_error
wg_device_rx_transport(struct wg_device *dev, struct wg_msg_transport *msg,
    size_t len, struct wg_session **s)
{
	struct wg_session *session;
	enum wg_error ret = WG_OK;
	size_t data_len = len - offsetof(struct wg_msg_transport, data);
	uint64_t counter = letoh64(msg->counter);

	if ((session = fm_lookup(&dev->d_sessions, msg->receiver)) == NULL)
		return WG_ID;

	wg_session_promote(session);

	/* TODO fix locks, at the moment we just kinda don't care */
	if (session->s_state != WG_STATE_INITIATOR &&
	    session->s_state != WG_STATE_RESPONDER)
		ret_error(WG_STATE);

	if (wg_timespec_timedout(&session->s_created, WG_REJECT_AFTER_TIME) ||
	    session->s_keyset.k_rxcounter > WG_REJECT_AFTER_MESSAGES)
		ret_error(WG_REJECT);

	if (!chacha20poly1305_decrypt(msg->data, msg->data, data_len, NULL, 0,
	    msg->counter, session->s_keyset.k_rxkey.k))
		ret_error(WG_DECRYPT);

	if (antireplay_update(&session->s_keyset.k_ar, counter))
		ret_error(WG_REPLAY);

	session->s_keyset.k_rxcounter = counter;
	session->s_peer->p_rx_bytes += data_len - WG_MAC_SIZE;

	if (session->s_state == WG_STATE_RESPONDER &&
	    wg_timespec_timedout(&session->s_created,
	    WG_REKEY_AFTER_TIME_RECV) &&
	    wg_timespec_timedout(&session->s_peer->p_last_initiation,
	    WG_REKEY_TIMEOUT))
		dev->d_outq(session->s_peer, WG_PKT_INITIATION, session->s_peer->p_id);

	*s = session;
leave:
	if (ret != WG_OK)
		wg_session_put(session);
	return ret;
}


enum wg_error
wg_device_tx_initiation(struct wg_device *dev, struct wg_msg_initiation *init,
    uint32_t id, struct wg_session **s)
{
	struct wg_peer *peer;
	struct wg_keypair kp;
	struct wg_handshake hs;
	struct wg_session *session = NULL;

	enum wg_error ret = WG_OK;

	if ((peer = fm_lookup(&dev->d_peers, id)) == NULL)
		return WG_ID;

	mtx_enter(&peer->p_mtx);
	if (!wg_timespec_timedout(&peer->p_last_initiation, WG_REKEY_TIMEOUT))
		ret = WG_HS_RATE;
	if (peer->p_attempts >= WG_REKEY_ATTEMPT_COUNT)
		ret = WG_HS_ATTEMPTS;
	if (ret == WG_OK) {
		getnanotime(&peer->p_last_initiation);
		peer->p_attempts++;
	}
	mtx_leave(&peer->p_mtx);

	if (ret != WG_OK)
		goto leave;

	/* We need to generate the session here first, so we can use s_local_id
	 * below. We also want to operate on a local handshake, so we don't
	 * have to lock the session. */
	session = wg_peer_new_session(peer);

	wg_device_getkey(dev, &kp);
	wg_keypair_generate(&hs.h_local);

	/* Noise handshake */
	init->type = WG_MSG_INITIATION;
	init->sender = session->s_local_id;
	memcpy(init->ephemeral, hs.h_local.pub.k, WG_KEY_SIZE);

	wg_hash2(hs.h_ck, WG_CONSTRUCTION, strlen(WG_CONSTRUCTION), NULL, 0);
	memcpy(hs.h_hash, hs.h_ck, WG_HASH_SIZE);
	wg_mix_hash(&hs, WG_IDENTIFIER, strlen(WG_IDENTIFIER));

	wg_mix_hash(&hs, peer->p_remote.k, WG_KEY_SIZE);
	wg_kdf(hs.h_ck, NULL, NULL, hs.h_ck, hs.h_local.pub.k, WG_KEY_SIZE);
	wg_mix_hash(&hs, hs.h_local.pub.k, WG_KEY_SIZE);
	wg_mix_dh(&hs, hs.h_local.priv.k, peer->p_remote.k);

	wg_handshake_encrypt(&hs, init->static_pub, kp.pub.k,
	    WG_KEY_SIZE);

	wg_mix_hash(&hs, init->static_pub, WG_ENCRYPTED_SIZE(WG_KEY_SIZE));
	wg_mix_dh(&hs, kp.priv.k, peer->p_remote.k);

	wg_timestamp_get(init->timestamp);

	wg_handshake_encrypt(&hs, init->timestamp, init->timestamp, WG_TIMESTAMP_SIZE);
	wg_mix_hash(&hs, init->timestamp, WG_ENCRYPTED_SIZE(WG_TIMESTAMP_SIZE));
	wg_hash2(hs.h_k, WG_MAC1, strlen(WG_MAC1), peer->p_remote.k, WG_KEY_SIZE);

	blake2s(init->mac1, (void *) init, hs.h_k, sizeof(init->mac1),
	    offsetof(struct wg_msg_initiation, mac1), sizeof(hs.h_k));
	memcpy(hs.h_mac, init->mac1, sizeof(hs.h_mac));

	/* TODO lock for cookie time? */
	if (wg_timespec_timedout(&peer->p_cookie.time, WG_COOKIE_VALID_TIME))
		bzero(init->mac2, WG_MAC_SIZE);
	else
		blake2s(init->mac2, (void *)init, peer->p_cookie.cookie,
		    sizeof(init->mac2), offsetof(struct wg_msg_initiation, mac2),
		    sizeof(peer->p_cookie.cookie));

	mtx_enter(&session->s_mtx);
	session->s_handshake = hs;
	session->s_state = WG_STATE_MADE_INITIATION;
	mtx_leave(&session->s_mtx);

	*s = session;
leave:
	/* if (ret != WG_OK)
		wg_session_put(session); */
	wg_peer_put(peer);
	explicit_bzero(&kp, sizeof(kp));
	explicit_bzero(&hs, sizeof(hs));
	return ret;
}

enum wg_error
wg_device_tx_response(struct wg_device *dev, struct wg_msg_response *resp,
    uint32_t id, struct wg_session **s)
{
	struct wg_handshake hs;
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((session = fm_lookup(&dev->d_sessions, id)) == NULL)
		return WG_ID;

	resp->type = WG_MSG_RESPONSE;
	resp->sender = session->s_local_id;
	resp->receiver = session->s_remote_id;

	mtx_enter(&session->s_mtx);
	hs = session->s_handshake;
	mtx_leave(&session->s_mtx);

	/* Noise handshake */
	wg_kdf(hs.h_ck, NULL, NULL, hs.h_ck, hs.h_local.pub.k, WG_KEY_SIZE);
	wg_mix_hash(&hs, hs.h_local.pub.k, WG_KEY_SIZE);

	memcpy(resp->ephemeral, hs.h_local.pub.k, WG_KEY_SIZE);

	wg_mix_dh(&hs, hs.h_local.priv.k, hs.h_remote.k);
	wg_mix_dh(&hs, hs.h_local.priv.k, session->s_peer->p_remote.k);

	wg_mix_psk(&hs, session->s_peer->p_shared.k);

	wg_handshake_encrypt(&hs, resp->empty, NULL, 0);

	wg_mix_hash(&hs, resp->empty, WG_ENCRYPTED_SIZE(0));

	wg_hash2(hs.h_k, WG_MAC1, strlen(WG_MAC1), session->s_peer->p_remote.k, WG_KEY_SIZE);
	blake2s(resp->mac1, (void *)resp, hs.h_k, sizeof(resp->mac1),
	    offsetof(struct wg_msg_response, mac1), sizeof(hs.h_k));
	memcpy(hs.h_mac, resp->mac1, sizeof(hs.h_mac));

	/* TODO lock for cookie time? */
	if (wg_timespec_timedout(&session->s_peer->p_cookie.time, WG_COOKIE_VALID_TIME))
		bzero(resp->mac2, WG_MAC_SIZE);
	else
		blake2s(resp->mac2, (void *)resp, session->s_peer->p_cookie.cookie,
		    sizeof(resp->mac2), offsetof(struct wg_msg_response, mac2),
		    sizeof(session->s_peer->p_cookie.cookie));

	mtx_enter(&session->s_mtx);
	if (session->s_state == WG_STATE_RECV_INITIATION)
		session->s_state = WG_STATE_MADE_RESPONSE;
	mtx_leave(&session->s_mtx);

	if (session->s_state != WG_STATE_MADE_RESPONSE)
		session->s_state = WG_STATE_MADE_RESPONSE;
		ret_error(WG_STATE);

	*s = session;
leave:
	if (ret != WG_OK)
		wg_session_put(session);
	return ret;
}

enum wg_error
wg_device_tx_cookie(struct wg_device *dev, struct wg_cookie *c,
    uint32_t sender, uint8_t mac[WG_MAC_SIZE], struct wg_msg_cookie *msg)
{
	uint8_t key[WG_KEY_SIZE]; // Same as WG_HASH_SIZE

	msg->type = WG_MSG_COOKIE;
	msg->receiver = sender;
	arc4random_buf(msg->nonce, sizeof(msg->nonce));

	wg_hash2(key, WG_COOKIE, strlen(WG_COOKIE), dev->d_keypair.pub.k,
	    WG_KEY_SIZE);
	xchacha20poly1305_encrypt(msg->value, c->cookie, WG_MAC_SIZE, mac,
	    WG_MAC_SIZE, msg->nonce, key);

	explicit_bzero(key, sizeof(key));
	return WG_OK;
}

enum wg_error
wg_device_tx_transport(struct wg_device *dev, struct wg_msg_transport *msg,
    size_t len, uint32_t id, struct wg_session **s)
{
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((session = fm_lookup(&dev->d_sessions, id)) == NULL)
		return WG_ID;

	/* TODO we should do some locking */
	if (session->s_state != WG_STATE_INITIATOR &&
	    session->s_state != WG_STATE_RESPONDER)
		ret_error(WG_STATE);

	if (wg_timespec_timedout(&session->s_created, WG_REJECT_AFTER_TIME) ||
	    session->s_keyset.k_txcounter > WG_REJECT_AFTER_MESSAGES)
		ret_error(WG_REJECT);

	msg->type = WG_MSG_TRANSPORT;
	msg->receiver = session->s_remote_id;
	msg->counter = htole64(session->s_keyset.k_txcounter++);

	chacha20poly1305_encrypt(msg->data, msg->data, len, NULL, 0, msg->counter,
	    session->s_keyset.k_txkey.k);

	/* Packet OK, but we do want a rekey */
	if (((session->s_state == WG_STATE_INITIATOR &&
	    wg_timespec_timedout(&session->s_created, WG_REKEY_AFTER_TIME)) ||
	    session->s_keyset.k_txcounter > WG_REKEY_AFTER_MESSAGES) &&
	    wg_timespec_timedout(&session->s_peer->p_last_initiation,
	    WG_REKEY_TIMEOUT))
		dev->d_outq(session->s_peer, WG_PKT_INITIATION, session->s_peer->p_id);

	session->s_peer->p_tx_bytes += len;
	*s = session;
leave:
	if (ret != WG_OK)
		wg_session_put(session);
	return ret;
}

/* WireGuard crypto helper functions */
void
wg_kdf(uint8_t first[WG_HASH_SIZE], uint8_t second[WG_HASH_SIZE],
    uint8_t third[WG_HASH_SIZE], uint8_t key[WG_KEY_SIZE], uint8_t * input,
    size_t input_len)
{
	uint8_t 	buffer[WG_HASH_SIZE + 1];
	uint8_t 	secret[WG_HASH_SIZE];

	blake2s_hmac(secret, input, key, WG_HASH_SIZE, input_len, WG_KEY_SIZE);

	if (!first)
		return;

	buffer[0] = 1;
	blake2s_hmac(buffer, buffer, secret, WG_HASH_SIZE, 1, WG_KEY_SIZE);
	memcpy(first, buffer, WG_HASH_SIZE);

	if (!second)
		return;

	buffer[WG_HASH_SIZE] = 2;
	blake2s_hmac(buffer, buffer, secret, WG_HASH_SIZE, WG_HASH_SIZE + 1, WG_KEY_SIZE);
	memcpy(second, buffer, WG_HASH_SIZE);

	if (!third)
		return;

	buffer[WG_HASH_SIZE] = 3;
	blake2s_hmac(buffer, buffer, secret, WG_HASH_SIZE, WG_HASH_SIZE + 1, WG_KEY_SIZE);
	memcpy(third, buffer, WG_HASH_SIZE);
}

void
wg_hash2(uint8_t out[WG_HASH_SIZE], uint8_t * in1, size_t in1_len, uint8_t * in2, size_t in2_len)
{
	struct blake2s_state s;
	blake2s_init(&s, WG_HASH_SIZE);
	blake2s_update(&s, in1, in1_len);
	blake2s_update(&s, in2, in2_len);
	blake2s_final(&s, out, WG_HASH_SIZE);
}

void
wg_mix_hash(struct wg_handshake * hs, uint8_t * in, size_t in_len)
{
	wg_hash2(hs->h_hash, hs->h_hash, WG_HASH_SIZE, in, in_len);
}

void
wg_mix_dh(struct wg_handshake * hs, uint8_t priv[WG_KEY_SIZE], uint8_t pub[WG_KEY_SIZE])
{
	uint8_t 	DH    [WG_KEY_SIZE];
	crypto_scalarmult_curve25519(DH, priv, pub);
	wg_kdf(hs->h_ck, hs->h_k, NULL, hs->h_ck, DH, WG_KEY_SIZE);
	explicit_bzero(DH, WG_KEY_SIZE);
}

void
wg_mix_psk(struct wg_handshake * hs, uint8_t psk[WG_KEY_SIZE])
{
	uint8_t 	t     [WG_HASH_SIZE];
	wg_kdf(hs->h_ck, t, hs->h_k, hs->h_ck, psk, WG_KEY_SIZE);
	wg_mix_hash(hs, t, WG_HASH_SIZE);
	explicit_bzero(t, WG_HASH_SIZE);
}

void
wg_handshake_encrypt(struct wg_handshake *hs, uint8_t *dst, uint8_t *src, size_t srclen)
{
	return chacha20poly1305_encrypt(dst, src, srclen, hs->h_hash, WG_HASH_SIZE, 0, hs->h_k);
}

int
wg_handshake_decrypt(struct wg_handshake *hs, uint8_t *dst, uint8_t *src, size_t srclen)
{
	return chacha20poly1305_decrypt(dst, src, srclen, hs->h_hash, WG_HASH_SIZE, 0, hs->h_k);
}

void
wg_timestamp_get(uint8_t ts[WG_TIMESTAMP_SIZE])
{
	struct timespec tv;
	getnanotime(&tv);
	*(uint64_t *) (ts) = htobe64(0x400000000000000aULL + tv.tv_sec);
	*(uint32_t *) (ts + 8) = htobe32(tv.tv_nsec);
}

int
wg_timespec_timedout(struct timespec * start, time_t timeout)
{
	struct timespec now;

	getnanotime(&now);

	return now.tv_sec == start->tv_sec + timeout ?
		now.tv_nsec > start->tv_nsec :
		now.tv_sec > start->tv_sec + timeout;
}

enum wg_pkt_type
wg_pkt_type(uint8_t *buf, size_t len)
{
	struct wg_msg_unknown *msg;

	if (len < sizeof(*msg))
		return WG_PKT_UNKNOWN;
	else
		msg = (struct wg_msg_unknown *)buf;

	if (msg->type == WG_MSG_INITIATION && len == wg_pkt_len[msg->type])
		return WG_PKT_INITIATION;
	else if (msg->type == WG_MSG_RESPONSE && len == wg_pkt_len[msg->type])
		return WG_PKT_RESPONSE;
	else if (msg->type == WG_MSG_COOKIE && len == wg_pkt_len[msg->type])
		return WG_PKT_COOKIE;
	else if (msg->type == WG_MSG_TRANSPORT && len >= wg_pkt_len[msg->type])
		return WG_PKT_TRANSPORT;
	else
		return WG_PKT_UNKNOWN;
}

void
wg_keypair_from_key(struct wg_keypair *kp, const struct wg_privkey *key)
{
	const uint8_t basepoint[WG_KEY_SIZE] = {9};

	/* memmove as key may overlap with kp->priv */
	memmove(kp->priv.k, key->k, WG_KEY_SIZE);

	/* We don't care if the input is not a valid key, we just set
	 * the bits and be done with it. The curve25519 library *SHOULD*
	 * do this, but since this may be returned to the user, we do
	 * it here too. */
	kp->priv.k[0] &= 248;
	kp->priv.k[31] &= 127;
	kp->priv.k[31] |= 64;

	crypto_scalarmult_curve25519(kp->pub.k, kp->priv.k, basepoint);
}

void
wg_keypair_generate(struct wg_keypair *kp)
{
	arc4random_buf(kp->priv.k, sizeof(kp->priv.k));
	wg_keypair_from_key(kp, &kp->priv);
}

enum wg_error
wg_msg_initiation_valid_mac2(struct wg_msg_initiation *msg, struct wg_cookie *c)
{
	uint8_t mac[WG_MAC_SIZE];
	blake2s(mac, (uint8_t *)msg, c->cookie, WG_MAC_SIZE,
	    offsetof(struct wg_msg_initiation, mac2), WG_COOKIE_SIZE);
	return timingsafe_bcmp(mac, msg->mac2, WG_MAC_SIZE) ? WG_MAC : WG_OK;
}

enum wg_error
wg_msg_response_valid_mac2(struct wg_msg_response *msg, struct wg_cookie *c)
{
	uint8_t mac[WG_MAC_SIZE];
	blake2s(mac, (uint8_t *)msg, c->cookie, WG_MAC_SIZE,
	    offsetof(struct wg_msg_response, mac2), WG_COOKIE_SIZE);
	return timingsafe_bcmp(mac, msg->mac2, WG_MAC_SIZE) ? WG_MAC : WG_OK;
}

void
wg_cookie_from_token(struct wg_cookie *c, struct wg_cookie_maker *cm,
    uint8_t *ip, uint8_t ip_len)
{
	if (wg_timespec_timedout(&cm->time, WG_COOKIE_VALID_TIME)) {
		getnanotime(&cm->time);
		arc4random_buf(cm->seed, sizeof(cm->seed));
	}

	blake2s(c->cookie, ip, cm->seed, WG_MAC_SIZE, ip_len, WG_COOKIE_SIZE);
}

/* Timer Functions */
void
wg_timer_setup(struct wg_timers *t, void *p, void (*keepalive)(void *),
    void (*broken)(void *), void (*reinit)(void *), void (*cleanup)(void *))
{
	assert(p != NULL);
	assert(keepalive != NULL);
	assert(broken != NULL);
	assert(reinit != NULL);
	assert(cleanup != NULL);

	timeout_set_proc(&t->t_ka, keepalive, p);
	timeout_set_proc(&t->t_pka, keepalive, p);
	timeout_set_proc(&t->t_broken, broken, p);
	timeout_set_proc(&t->t_reinit, reinit, p);
	timeout_set_proc(&t->t_cleanup, cleanup, p);
}

void
wg_timer_stop(struct wg_timers *t)
{
	/* TODO need barrier? */
	timeout_del(&t->t_ka);
	timeout_del(&t->t_pka);
	timeout_del(&t->t_broken);
	timeout_del(&t->t_reinit);
	timeout_del(&t->t_cleanup);
}

void
wg_timer_persistent_keepalive_tick(struct wg_timers *t)
{
	if (t->t_pka_interval > 0)
		timeout_add_sec(&t->t_pka, t->t_pka_interval);
}

uint16_t
wg_timer_persistent_keepalive_get(struct wg_timers *t)
{
	return t->t_pka_interval;
}

void
wg_timer_persistent_keepalive_set(struct wg_timers *t, uint16_t interval)
{
	t->t_pka_interval = interval;
}

void
wg_timer_cleanup_tick(struct wg_timers *t)
{
	timeout_add_sec(&t->t_cleanup, WG_REJECT_AFTER_TIME * 3);
}

void
wg_timer_keepalive_flag(struct wg_timers *t)
{
	timeout_add_sec(&t->t_ka, WG_KEEPALIVE_TIMEOUT);
}

void
wg_timer_keepalive_unflag(struct wg_timers *t)
{
	timeout_del(&t->t_ka);
}

void
wg_timer_broken_flag(struct wg_timers *t)
{
	if (timeout_pending(&t->t_broken) == 0)
		timeout_add_sec(&t->t_broken, WG_REKEY_TIMEOUT + WG_KEEPALIVE_TIMEOUT);
}

void
wg_timer_broken_unflag(struct wg_timers *t)
{
	timeout_del(&t->t_broken);
}

void
wg_timer_reinit_flag(struct wg_timers *t)
{
	timeout_add_sec(&t->t_reinit, WG_REKEY_TIMEOUT);
}

void
wg_timer_reinit_unflag(struct wg_timers *t)
{
	timeout_del(&t->t_reinit);
}