4 files changed, 88 insertions, 90 deletions

diff --git a/src/cipher.c b/src/cipher.c
index 0ca48b4..22fba9a 100644
--- a/src/cipher.c
+++ b/src/cipher.c
@@ -206,36 +206,38 @@ buffer_t cipher_aesctr_calc_ctr(cipher_t* c, key_derivation_t* kd, seq_nr_t seq_
   if(!c->salt_.buf_) {
     c->salt_.length_ = 14;
     c->salt_.buf_ = malloc(c->salt_.length_);
-    if(c->salt_.buf_)
+    if(!c->salt_.buf_)
       return result;
   }
   int ret = key_derivation_generate(kd, LABEL_SATP_SALT, seq_nr, c->salt_.buf_, c->salt_.length_);
-  if(ret)
+  if(ret < 0)
     return result;
 
   mpz_t ctr, sid_mux, seq;
   mpz_init2(ctr, 128);
-  mpz_init2(sid_mux, 96);
-  mpz_init2(seq, 48);
+  mpz_init2(sid_mux, 128);
+  mpz_init2(seq, 128);
   
   int faked_msb = 0;
-  if(!c->salt_.buf_[0])
+  if(!c->salt_.buf_[0]) {
     c->salt_.buf_[0] = 1;
+    faked_msb = 1;
+  }
 
   mpz_import(ctr, c->salt_.length_, 1, 1, 0, 0, c->salt_.buf_);
-  mpz_mul_2exp(ctr, ctr, 16);
 
   mpz_set_ui(sid_mux, mux);
-  mpz_mul_2exp(sid_mux, sid_mux, 16);
+  mpz_mul_2exp(sid_mux, sid_mux, (sizeof(sender_id) * 8));
   mpz_add_ui(sid_mux, sid_mux, sender_id);
-  mpz_mul_2exp(sid_mux, sid_mux, 64);
+  mpz_mul_2exp(sid_mux, sid_mux, 48);
 
   mpz_set_ui(seq, seq_nr);
-  mpz_mul_2exp(seq, seq, 16);
 
   mpz_xor(ctr, ctr, sid_mux);
   mpz_xor(ctr, ctr, seq);
 
+  mpz_mul_2exp(ctr, ctr, 16);
+
   result.buf_ = mpz_export(NULL, (size_t*)&result.length_, 1, 1, 0, 0, ctr);
   if(faked_msb) {
     c->salt_.buf_[0] = 0;
@@ -251,8 +253,8 @@ buffer_t cipher_aesctr_calc_ctr(cipher_t* c, key_derivation_t* kd, seq_nr_t seq_
 
 u_int32_t cipher_aesctr_crypt(cipher_t* c, key_derivation_t* kd, u_int8_t* in, u_int32_t ilen, u_int8_t* out, u_int32_t olen, seq_nr_t seq_nr, sender_id_t sender_id, mux_t mux)
 {
-  if(!c || !c->key_.buf_ || !c->salt_.buf_) {
-    log_printf(ERR, "cipher not initialized or no key or salt set");
+  if(!c) {
+    log_printf(ERR, "cipher not initialized");
     return 0;
   }
   if(!kd) {
@@ -269,7 +271,7 @@ u_int32_t cipher_aesctr_crypt(cipher_t* c, key_derivation_t* kd, u_int8_t* in, u
   }
 
   int ret = key_derivation_generate(kd, LABEL_SATP_ENCRYPTION, seq_nr, c->key_.buf_, c->key_.length_);
-  if(ret)
+  if(ret < 0)
     return 0;
   
   gcry_error_t err = gcry_cipher_setkey(c->handle_, c->key_.buf_, c->key_.length_);
diff --git a/src/key_derivation.c b/src/key_derivation.c
index 44ae667..3977993 100644
--- a/src/key_derivation.c
+++ b/src/key_derivation.c
@@ -60,15 +60,17 @@ int key_derivation_init(key_derivation_t* kd, const char* type, int8_t ld_kdr, u
   }
 
   kd->ld_kdr_ = ld_kdr;
-  if(ld_kdr > (sizeof(seq_nr_t) * 8))
+  if(ld_kdr > (int8_t)(sizeof(seq_nr_t) * 8))
     kd->ld_kdr_ = sizeof(seq_nr_t) * 8;
+
   kd->key_length_ = key_len * sizeof(key[0]) * 8;
   kd->handle_ = 0;
 
   int i;
   for(i = 0; i<KD_LABEL_COUNT; ++i) {
-    kd->key_store_[i].buf_ = NULL;
-    kd->key_store_[i].length_ = 0;
+    kd->key_store_[i].key_.buf_ = NULL;
+    kd->key_store_[i].key_.length_ = 0;
+    kd->key_store_[i].r_ = 0;
   }
 
   if(!key) {
@@ -123,8 +125,8 @@ void key_derivation_close(key_derivation_t* kd)
 
   int i;
   for(i = 0; i<KD_LABEL_COUNT; ++i) {
-    if(kd->key_store_[i].buf_)
-      free(kd->key_store_[i].buf_);
+    if(kd->key_store_[i].key_.buf_)
+      free(kd->key_store_[i].key_.buf_);
   }
 }
 
@@ -200,75 +202,50 @@ void key_derivation_aesctr_close(key_derivation_t* kd)
     gcry_cipher_close(kd->handle_);
 }
 
-int key_derivation_aesctr_calc_ctr(key_derivation_t* kd, buffer_t* result, satp_prf_label_t label, seq_nr_t seq_nr)
+int key_derivation_aesctr_calc_ctr(key_derivation_t* kd, key_store_t* result, satp_prf_label_t label, seq_nr_t seq_nr)
 {
   if(!kd || !result)
     return -1;
 
-  result->buf_ = NULL;
-  result->length_ = 0;
+  seq_nr_t r = 0;
+  if(kd->ld_kdr_ >= 0)
+    r = seq_nr >> kd->ld_kdr_;
 
-  mpz_t ctr, key_id, r, seq;
+  if(kd->key_store_[label].key_.buf_ && kd->key_store_[label].r_ == r) {
+    if(!r || (seq_nr % r))
+      return 0;
+  }
+  result->r_ = r;
+
+  mpz_t ctr, key_id;
   mpz_init2(ctr, 128);
   mpz_init2(key_id, 128);
-  mpz_init2(r, 128);
-  mpz_init2(seq, (sizeof(seq_nr_t) * 8));
-
-  mpz_set_ui(seq, seq_nr);
 
   int faked_msb = 0;
-  if(!kd->master_salt_.buf_[0])
+  if(!kd->master_salt_.buf_[0]) {
     kd->master_salt_.buf_[0] = 1;
-
-  if(kd->ld_kdr_ == -1)
-    mpz_set_ui(r, 0);
-  else
-    mpz_fdiv_q_2exp(r, seq, kd->ld_kdr_);
-
-  if(kd->key_store_[label].buf_) {
-    if(!mpz_cmp_ui(r, 0)) {
-      mpz_clear(seq);
-      mpz_clear(ctr);
-      mpz_clear(key_id);
-      mpz_clear(r);
-      return 0;
-    }
-    
-    mpz_t mod;
-    mpz_init2(mod, (sizeof(seq_nr_t) * 8));
-    mpz_fdiv_r(mod, seq, r);
-    if(mpz_cmp_ui(mod, 0)) {
-      mpz_clear(seq);
-      mpz_clear(mod);
-      mpz_clear(ctr);
-      mpz_clear(key_id);
-      mpz_clear(r);
-      return 0;
-    }
-    mpz_clear(mod);
+    faked_msb = 1;
   }
+  mpz_import(ctr, kd->master_salt_.length_, 1, 1, 0, 0, kd->master_salt_.buf_);
 
   mpz_set_ui(key_id, label);
-  mpz_mul_2exp(key_id, key_id, (sizeof(seq_nr_t) * 8));
-  mpz_add(key_id, key_id, r);
-
-  mpz_import(ctr, kd->master_salt_.length_, 1, 1, 0, 0, kd->master_salt_.buf_);
+  mpz_mul_2exp(key_id, key_id, (sizeof(r) * 8));
+  mpz_add_ui(key_id, key_id, r);
 
   mpz_xor(ctr, ctr, key_id);
   mpz_mul_2exp(ctr, ctr, 16);
 
-  if(result->buf_)
-    free(result->buf_);
-  result->buf_ = mpz_export(NULL, (size_t*)&(result->length_), 1, 1, 0, 0, ctr);
+  if(result->key_.buf_)
+    free(result->key_.buf_);
+  result->key_.buf_ = mpz_export(NULL, (size_t*)&(result->key_.length_), 1, 1, 0, 0, ctr);
+
   if(faked_msb) {
     kd->master_salt_.buf_[0] = 0;
-    result->buf_[0] = 0;
+    result->key_.buf_[0] = 0;
   }
 
-  mpz_clear(seq);
   mpz_clear(ctr);
   mpz_clear(key_id);
-  mpz_clear(r);
 
   return 1;
 }
@@ -280,21 +257,22 @@ int key_derivation_aesctr_generate(key_derivation_t* kd, satp_prf_label_t label,
     return -1;
   }
 
-  buffer_t ctr;
-  ctr.buf_ = NULL;
-  ctr.length_ = 0;
+  key_store_t ctr;
+  ctr.key_.buf_ = NULL;
+  ctr.key_.length_ = 0;
+  ctr.r_ = 0;
   int ret = key_derivation_aesctr_calc_ctr(kd, &ctr, label, seq_nr);
   if(ret < 0) {
     log_printf(ERR, "failed to calculate key derivation CTR");
     return -1;
   }
   else if(!ret) {
-    if(len > kd->key_store_[label].length_) {
+    if(len > kd->key_store_[label].key_.length_) {
       log_printf(WARNING, "stored (old) key for label 0x%02X is too short, filling with zeros", label);
       memset(key, 0, len);
-      len = kd->key_store_[label].length_;
+      len = kd->key_store_[label].key_.length_;
     }
-    memcpy(key, kd->key_store_[label].buf_, len);
+    memcpy(key, kd->key_store_[label].key_.buf_, len);
     return 0;
   }
 
@@ -304,8 +282,8 @@ int key_derivation_aesctr_generate(key_derivation_t* kd, satp_prf_label_t label,
     return -1;
   }
 
-  err = gcry_cipher_setctr(kd->handle_, ctr.buf_, ctr.length_);
-  free(ctr.buf_);
+  err = gcry_cipher_setctr(kd->handle_, ctr.key_.buf_, ctr.key_.length_);
+  free(ctr.key_.buf_);
 
   if(err) {
     log_printf(ERR, "failed to set key derivation CTR: %s/%s", gcry_strerror(err), gcry_strsource(err));
@@ -319,26 +297,30 @@ int key_derivation_aesctr_generate(key_derivation_t* kd, satp_prf_label_t label,
     return -1;
   }
   
-  if(!kd->key_store_[label].buf_) {
-    kd->key_store_[label].length_ = 0;
-    kd->key_store_[label].buf_ = malloc(len);
-    if(!kd->key_store_[label].buf_) {
+  if(!kd->ld_kdr_)
+    return 1;
+
+  if(!kd->key_store_[label].key_.buf_) {
+    kd->key_store_[label].key_.length_ = 0;
+    kd->key_store_[label].key_.buf_ = malloc(len);
+    if(!kd->key_store_[label].key_.buf_) {
       log_printf(ERR, "memory error at key derivation");
       return -2;
     }
-    kd->key_store_[label].length_ = len;
+    kd->key_store_[label].key_.length_ = len;
   }
-  else if(kd->key_store_[label].length_ < len) {
-    u_int8_t* tmp = realloc(kd->key_store_[label].buf_, len);
+  else if(kd->key_store_[label].key_.length_ < len) {
+    u_int8_t* tmp = realloc(kd->key_store_[label].key_.buf_, len);
     if(!tmp) {
       log_printf(ERR, "memory error at key derivation");
       return -2;
     }
-    kd->key_store_[label].buf_ = tmp;
-    kd->key_store_[label].length_ = len;
+    kd->key_store_[label].key_.buf_ = tmp;
+    kd->key_store_[label].key_.length_ = len;
   }
 
-  memcpy(kd->key_store_[label].buf_, key, len);
+  memcpy(kd->key_store_[label].key_.buf_, key, len);
+  kd->key_store_[label].r_ = ctr.r_;
 
-  return 0;
+  return 1;
 }
diff --git a/src/key_derivation.h b/src/key_derivation.h
index 21f069d..d8f49c6 100644
--- a/src/key_derivation.h
+++ b/src/key_derivation.h
@@ -48,6 +48,12 @@ typedef enum satp_prf_label_enum satp_prf_label_t;
 enum key_derivation_type_enum { kd_unknown, kd_null, kd_aes_ctr };
 typedef enum key_derivation_type_enum key_derivation_type_t;
 
+struct key_store_struct {
+  buffer_t key_;
+  seq_nr_t r_;
+};
+typedef struct key_store_struct key_store_t;
+
 struct key_derivation_struct {
   key_derivation_type_t type_;
   int8_t ld_kdr_;
@@ -55,7 +61,7 @@ struct key_derivation_struct {
   buffer_t master_key_;
   buffer_t master_salt_;
   gcry_cipher_hd_t handle_;
-  buffer_t key_store_[KD_LABEL_COUNT];
+  key_store_t key_store_[KD_LABEL_COUNT];
 };
 typedef struct key_derivation_struct key_derivation_t;
 
@@ -67,7 +73,7 @@ int key_derivation_null_generate(u_int8_t* key, u_int32_t len);
 
 int key_derivation_aesctr_init(key_derivation_t* kd, u_int16_t key_length);
 void key_derivation_aesctr_close(key_derivation_t* kd);
-int key_derivation_aesctr_calc_ctr(key_derivation_t* kd, buffer_t* result, satp_prf_label_t label, seq_nr_t seq_nr);
+int key_derivation_aesctr_calc_ctr(key_derivation_t* kd, key_store_t* result, satp_prf_label_t label, seq_nr_t seq_nr);
 int key_derivation_aesctr_generate(key_derivation_t* kd, satp_prf_label_t label, seq_nr_t seq_nr, u_int8_t* key, u_int32_t len);
 
 #endif
diff --git a/src/uanytun.c b/src/uanytun.c
index 8f677a4..8ad6ed7 100644
--- a/src/uanytun.c
+++ b/src/uanytun.c
@@ -105,10 +105,17 @@ int main_loop(tun_device_t* dev, udp_socket_t* sock, options_t* opt)
     return_value = ret;
   }
 
-  key_derivation_t kd;
-  ret = key_derivation_init(&kd, opt->kd_prf_, opt->ld_kdr_, opt->key_.buf_, opt->key_.length_, opt->salt_.buf_, opt->salt_.length_);
+  key_derivation_t kd_in;
+  ret = key_derivation_init(&kd_in, opt->kd_prf_, opt->ld_kdr_, opt->key_.buf_, opt->key_.length_, opt->salt_.buf_, opt->salt_.length_);
   if(ret) {
-    log_printf(ERR, "could not initialize key derivation of type %s", opt->kd_prf_);
+    log_printf(ERR, "could not initialize inbound key derivation of type %s", opt->kd_prf_);
+    return_value = ret;
+  }
+
+  key_derivation_t kd_out;
+  ret = key_derivation_init(&kd_out, opt->kd_prf_, opt->ld_kdr_, opt->key_.buf_, opt->key_.length_, opt->salt_.buf_, opt->salt_.length_);
+  if(ret) {
+    log_printf(ERR, "could not initialize outbound key derivation of type %s", opt->kd_prf_);
     return_value = ret;
   }
 
@@ -158,7 +165,7 @@ int main_loop(tun_device_t* dev, udp_socket_t* sock, options_t* opt)
       else
         plain_packet_set_type(&plain_packet, PAYLOAD_TYPE_UNKNOWN);
       
-      cipher_encrypt(&c, &kd, &plain_packet, &encrypted_packet, seq_nr, opt->sender_id_, opt->mux_); 
+      cipher_encrypt(&c, &kd_out, &plain_packet, &encrypted_packet, seq_nr, opt->sender_id_, opt->mux_); 
       seq_nr++;
       
           // TODO: add auth-tag
@@ -201,7 +208,7 @@ int main_loop(tun_device_t* dev, udp_socket_t* sock, options_t* opt)
         free(addrstring);
       }
       
-      cipher_decrypt(&c, &kd, &encrypted_packet, &plain_packet); 
+      cipher_decrypt(&c, &kd_in, &encrypted_packet, &plain_packet); 
       
       len = tun_write(dev, plain_packet_get_payload(&plain_packet), plain_packet_get_payload_length(&plain_packet));
       if(len == -1)
@@ -210,7 +217,8 @@ int main_loop(tun_device_t* dev, udp_socket_t* sock, options_t* opt)
   }
 
   cipher_close(&c);
-  key_derivation_close(&kd);
+  key_derivation_close(&kd_out);
+  key_derivation_close(&kd_in);
   seq_win_clear(&seq_win);
 
   return return_value;