1 files changed, 49 insertions, 6 deletions

diff --git a/media/mtransport/test_nr_socket.cpp b/media/mtransport/test_nr_socket.cpp
index e7dad79cd1..fb2de3096e 100644
--- a/media/mtransport/test_nr_socket.cpp
+++ b/media/mtransport/test_nr_socket.cpp
@@ -200,6 +200,7 @@ int TestNat::create_socket_factory(nr_socket_factory **factorypp) {
 
 TestNrSocket::TestNrSocket(TestNat *nat)
   : nat_(nat),
+    tls_(false),
     timer_handle_(nullptr) {
   nat_->insert_socket(this);
 }
@@ -473,6 +474,10 @@ int TestNrSocket::connect(nr_transport_addr *addr) {
     return R_INTERNAL;
   }
 
+  if (addr->tls_host[0] != '\0') {
+    tls_ = true;
+  }
+
   if (!nat_->enabled_
       || addr->protocol==IPPROTO_UDP  // Horrible hack to allow default address
                                       // discovery to work. Only works because
@@ -508,6 +513,24 @@ int TestNrSocket::connect(nr_transport_addr *addr) {
 }
 
 int TestNrSocket::write(const void *msg, size_t len, size_t *written) {
+  UCHAR *buf = static_cast<UCHAR*>(const_cast<void*>(msg));
+  if (nat_->block_stun_ && nr_is_stun_message(buf, len)) {
+    // Should cause this socket to be abandoned
+    r_log(LOG_GENERIC, LOG_DEBUG,
+          "TestNrSocket %s dropping outgoing TCP "
+          "because it is configured to drop STUN",
+          my_addr().as_string);
+    return R_INTERNAL;
+  }
+
+  if (nat_->block_tcp_ && !tls_) {
+    // Should cause this socket to be abandoned
+    r_log(LOG_GENERIC, LOG_DEBUG,
+          "TestNrSocket %s dropping outgoing TCP "
+          "because it is configured to drop TCP",
+          my_addr().as_string);
+    return R_INTERNAL;
+  }
 
   if (port_mappings_.empty()) {
     // The no-nat case, just pass call through.
@@ -518,7 +541,11 @@ int TestNrSocket::write(const void *msg, size_t len, size_t *written) {
   } else {
     destroy_stale_port_mappings();
     if (port_mappings_.empty()) {
-      return -1;
+      r_log(LOG_GENERIC, LOG_DEBUG,
+            "TestNrSocket %s dropping outgoing TCP "
+            "because the port mapping was stale",
+            my_addr().as_string);
+      return R_INTERNAL;
     }
     // This is TCP only
     MOZ_ASSERT(port_mappings_.size() == 1);
@@ -533,18 +560,34 @@ int TestNrSocket::write(const void *msg, size_t len, size_t *written) {
 }
 
 int TestNrSocket::read(void *buf, size_t maxlen, size_t *len) {
+  int r;
 
   if (port_mappings_.empty()) {
-    return internal_socket_->read(buf, maxlen, len);
+    r = internal_socket_->read(buf, maxlen, len);
   } else {
     MOZ_ASSERT(port_mappings_.size() == 1);
-    int bytesRead =
-      port_mappings_.front()->external_socket_->read(buf, maxlen, len);
-    if (bytesRead > 0 && nat_->refresh_on_ingress_) {
+    r = port_mappings_.front()->external_socket_->read(buf, maxlen, len);
+    if (!r && nat_->refresh_on_ingress_) {
       port_mappings_.front()->last_used_ = PR_IntervalNow();
     }
-    return bytesRead;
   }
+
+  if (r) {
+    return r;
+  }
+
+  if (nat_->block_tcp_ && !tls_) {
+    // Should cause this socket to be abandoned
+    return R_INTERNAL;
+  }
+
+  UCHAR *cbuf = static_cast<UCHAR*>(const_cast<void*>(buf));
+  if (nat_->block_stun_ && nr_is_stun_message(cbuf, *len)) {
+    // Should cause this socket to be abandoned
+    return R_INTERNAL;
+  }
+
+  return r;
 }
 
 int TestNrSocket::async_wait(int how, NR_async_cb cb, void *cb_arg,