1 files changed, 20 insertions, 6 deletions
diff --git a/security/nss/gtests/ssl_gtest/test_io.h b/security/nss/gtests/ssl_gtest/test_io.h
index ac2497222..dbeb6b9d4 100644
--- a/security/nss/gtests/ssl_gtest/test_io.h
+++ b/security/nss/gtests/ssl_gtest/test_io.h
@@ -33,9 +33,18 @@ class PacketFilter {
     CHANGE,  // change the packet to a different value
     DROP     // drop the packet
   };
-
+  PacketFilter(bool enabled = true) : enabled_(enabled) {}
   virtual ~PacketFilter() {}
 
+  virtual Action Process(const DataBuffer& input, DataBuffer* output) {
+    if (!enabled_) {
+      return KEEP;
+    }
+    return Filter(input, output);
+  }
+  void Enable() { enabled_ = true; }
+  void Disable() { enabled_ = false; }
+
   // The packet filter takes input and has the option of mutating it.
   //
   // A filter that modifies the data places the modified data in *output and
@@ -43,6 +52,9 @@ class PacketFilter {
   // case the value in *output is ignored.  A Filter can return DROP, in which
   // case the packet is dropped (and *output is ignored).
   virtual Action Filter(const DataBuffer& input, DataBuffer* output) = 0;
+
+ private:
+  bool enabled_;
 };
 
 class DummyPrSocket : public DummyIOLayerMethods {
@@ -53,7 +65,7 @@ class DummyPrSocket : public DummyIOLayerMethods {
         peer_(),
         input_(),
         filter_(nullptr),
-        writeable_(true) {}
+        write_error_(0) {}
   virtual ~DummyPrSocket() {}
 
   // Create a file descriptor that will reference this object.  The fd must not
@@ -62,7 +74,9 @@ class DummyPrSocket : public DummyIOLayerMethods {
 
   std::weak_ptr<DummyPrSocket>& peer() { return peer_; }
   void SetPeer(const std::shared_ptr<DummyPrSocket>& peer) { peer_ = peer; }
-  void SetPacketFilter(std::shared_ptr<PacketFilter> filter);
+  void SetPacketFilter(const std::shared_ptr<PacketFilter>& filter) {
+    filter_ = filter;
+  }
   // Drops peer, packet filter and any outstanding packets.
   void Reset();
 
@@ -71,7 +85,7 @@ class DummyPrSocket : public DummyIOLayerMethods {
   int32_t Recv(PRFileDesc* f, void* buf, int32_t buflen, int32_t flags,
                PRIntervalTime to) override;
   int32_t Write(PRFileDesc* f, const void* buf, int32_t length) override;
-  void CloseWrites() { writeable_ = false; }
+  void SetWriteError(PRErrorCode code) { write_error_ = code; }
 
   SSLProtocolVariant variant() const { return variant_; }
   bool readable() const { return !input_.empty(); }
@@ -98,7 +112,7 @@ class DummyPrSocket : public DummyIOLayerMethods {
   std::weak_ptr<DummyPrSocket> peer_;
   std::queue<Packet> input_;
   std::shared_ptr<PacketFilter> filter_;
-  bool writeable_;
+  PRErrorCode write_error_;
 };
 
 // Marker interface.
@@ -164,6 +178,6 @@ class Poller {
       timers_;
 };
 
-}  // end of namespace
+}  // namespace nss_test
 
 #endif