IPv4: Basic implementation of TCP socket shutdown

We can now participate in the TCP connection closing handshake. :^)
This implementation is definitely not complete and needs to handle a
bunch of other cases. But it's a huge improvement over not being able
to close connections at all.

Note that we hold on to pending-close sockets indefinitely, until they
are moved into the Closed state. This should also have a timeout but
that's still a FIXME. :^)

Fixes #428.

7 files changed, 69 insertions, 8 deletions

diff --git a/Kernel/Net/IPv4Socket.cpp b/Kernel/Net/IPv4Socket.cpp
index b19aba3b8f..35c8c472b0 100644
--- a/Kernel/Net/IPv4Socket.cpp
+++ b/Kernel/Net/IPv4Socket.cpp
@@ -500,3 +500,8 @@ int IPv4Socket::ioctl(FileDescription&, unsigned request, unsigned arg)
 
     return -EINVAL;
 }
+
+void IPv4Socket::close()
+{
+    shutdown(SHUT_RDWR);
+}
diff --git a/Kernel/Net/IPv4Socket.h b/Kernel/Net/IPv4Socket.h
index f00cc84cee..ebf09c38ae 100644
--- a/Kernel/Net/IPv4Socket.h
+++ b/Kernel/Net/IPv4Socket.h
@@ -46,6 +46,7 @@ public:
 
     static Lockable<HashTable<IPv4Socket*>>& all_sockets();
 
+    virtual void close() override;
     virtual KResult bind(const sockaddr*, socklen_t) override;
     virtual KResult connect(FileDescription&, const sockaddr*, socklen_t, ShouldBlock = ShouldBlock::Yes) override;
     virtual KResult listen(int) override;
diff --git a/Kernel/Net/NetworkTask.cpp b/Kernel/Net/NetworkTask.cpp
index 1c52923fcb..0d93f2bd0e 100644
--- a/Kernel/Net/NetworkTask.cpp
+++ b/Kernel/Net/NetworkTask.cpp
@@ -572,6 +572,9 @@ void handle_tcp(const IPv4Packet& ipv4_packet)
             socket->set_ack_number(tcp_packet.sequence_number() + payload_size + 1);
             socket->set_state(TCPSocket::State::TimeWait);
             return;
+        case TCPFlags::ACK | TCPFlags::RST:
+            socket->set_state(TCPSocket::State::Closed);
+            return;
         default:
             kprintf("handle_tcp: unexpected flags in FinWait2 state\n");
             socket->send_tcp_packet(TCPFlags::RST);
@@ -596,12 +599,8 @@ void handle_tcp(const IPv4Packet& ipv4_packet)
                 socket->did_receive(ipv4_packet.source(), tcp_packet.source_port(), KBuffer::copy(&ipv4_packet, sizeof(IPv4Packet) + ipv4_packet.payload_size()));
 
             socket->set_ack_number(tcp_packet.sequence_number() + payload_size + 1);
-            // TODO: We should only send a FIN packet out once we're shutting
-            // down our side of the socket, so we should change this back to
-            // just being an ACK and a transition to CloseWait once we have a
-            // shutdown() implementation.
-            socket->send_tcp_packet(TCPFlags::FIN | TCPFlags::ACK);
-            socket->set_state(TCPSocket::State::Closing);
+            socket->send_tcp_packet(TCPFlags::ACK);
+            socket->set_state(TCPSocket::State::CloseWait);
             socket->set_connected(false);
             return;
         }
diff --git a/Kernel/Net/Socket.cpp b/Kernel/Net/Socket.cpp
index 6e66db59c7..9adfd9278a 100644
--- a/Kernel/Net/Socket.cpp
+++ b/Kernel/Net/Socket.cpp
@@ -165,7 +165,13 @@ KResult Socket::shutdown(int how)
 {
     if (type() == SOCK_STREAM && !is_connected())
         return KResult(-ENOTCONN);
-    m_shut_down_for_reading |= how & SHUT_RD;
-    m_shut_down_for_writing |= how & SHUT_WR;
+    if (m_role == Role::Listener)
+        return KResult(-ENOTCONN);
+    if (!m_shut_down_for_writing && (how & SHUT_WR))
+        shut_down_for_writing();
+    if (!m_shut_down_for_reading && (how & SHUT_RD))
+        shut_down_for_reading();
+    m_shut_down_for_reading |= (how & SHUT_RD) != 0;
+    m_shut_down_for_writing |= (how & SHUT_WR) != 0;
     return KSuccess;
 }
diff --git a/Kernel/Net/Socket.h b/Kernel/Net/Socket.h
index a0a7e1d209..c74b0f8b3b 100644
--- a/Kernel/Net/Socket.h
+++ b/Kernel/Net/Socket.h
@@ -141,6 +141,9 @@ protected:
 
     virtual const char* class_name() const override { return "Socket"; }
 
+    virtual void shut_down_for_reading() {}
+    virtual void shut_down_for_writing() {}
+
     Role m_role { Role::None };
 
 protected:
diff --git a/Kernel/Net/TCPSocket.cpp b/Kernel/Net/TCPSocket.cpp
index 8e114ddfa5..618aa3797f 100644
--- a/Kernel/Net/TCPSocket.cpp
+++ b/Kernel/Net/TCPSocket.cpp
@@ -55,6 +55,19 @@ void TCPSocket::set_state(State new_state)
 
     if (new_state == State::Established && m_direction == Direction::Outgoing)
         m_role = Role::Connected;
+
+    if (new_state == State::Closed) {
+        LOCKER(closing_sockets().lock());
+        closing_sockets().resource().remove(tuple());
+    }
+}
+
+Lockable<HashMap<IPv4SocketTuple, RefPtr<TCPSocket>>>& TCPSocket::closing_sockets()
+{
+    static Lockable<HashMap<IPv4SocketTuple, RefPtr<TCPSocket>>>* s_map;
+    if (!s_map)
+        s_map = new Lockable<HashMap<IPv4SocketTuple, RefPtr<TCPSocket>>>;
+    return *s_map;
 }
 
 Lockable<HashMap<IPv4SocketTuple, TCPSocket*>>& TCPSocket::sockets_by_tuple()
@@ -137,6 +150,10 @@ TCPSocket::~TCPSocket()
 {
     LOCKER(sockets_by_tuple().lock());
     sockets_by_tuple().resource().remove(tuple());
+
+#ifdef TCP_SOCKET_DEBUG
+    dbg() << "~TCPSocket in state " << to_string(state());
+#endif
 }
 
 NonnullRefPtr<TCPSocket> TCPSocket::create(int protocol)
@@ -420,3 +437,27 @@ bool TCPSocket::protocol_is_disconnected() const
         return false;
     }
 }
+
+void TCPSocket::shut_down_for_writing()
+{
+    if (state() == State::Established) {
+        dbg() << " Sending FIN/ACK from Established and moving into FinWait1";
+        send_tcp_packet(TCPFlags::FIN | TCPFlags::ACK);
+        set_state(State::FinWait1);
+    } else {
+        dbg() << " Shutting down TCPSocket for writing but not moving to FinWait1 since state is " << to_string(state());
+    }
+}
+
+void TCPSocket::close()
+{
+    IPv4Socket::close();
+    if (state() == State::CloseWait) {
+        dbg() << " Sending FIN from CloseWait and moving into LastAck";
+        send_tcp_packet(TCPFlags::FIN | TCPFlags::ACK);
+        set_state(State::LastAck);
+    }
+
+    LOCKER(closing_sockets().lock());
+    closing_sockets().resource().set(tuple(), *this);
+}
diff --git a/Kernel/Net/TCPSocket.h b/Kernel/Net/TCPSocket.h
index 22cf796225..257c3804ac 100644
--- a/Kernel/Net/TCPSocket.h
+++ b/Kernel/Net/TCPSocket.h
@@ -155,12 +155,16 @@ public:
     static RefPtr<TCPSocket> from_tuple(const IPv4SocketTuple& tuple);
     static RefPtr<TCPSocket> from_endpoints(const IPv4Address& local_address, u16 local_port, const IPv4Address& peer_address, u16 peer_port);
 
+    static Lockable<HashMap<IPv4SocketTuple, RefPtr<TCPSocket>>>& closing_sockets();
+
     RefPtr<TCPSocket> create_client(const IPv4Address& local_address, u16 local_port, const IPv4Address& peer_address, u16 peer_port);
     void set_originator(TCPSocket& originator) { m_originator = originator.make_weak_ptr(); }
     bool has_originator() { return !!m_originator; }
     void release_to_originator();
     void release_for_accept(RefPtr<TCPSocket>);
 
+    virtual void close() override;
+
 protected:
     void set_direction(Direction direction) { m_direction = direction; }
 
@@ -170,6 +174,8 @@ private:
 
     static NetworkOrdered<u16> compute_tcp_checksum(const IPv4Address& source, const IPv4Address& destination, const TCPPacket&, u16 payload_size);
 
+    virtual void shut_down_for_writing() override;
+
     virtual int protocol_receive(const KBuffer&, void* buffer, size_t buffer_size, int flags) override;
     virtual int protocol_send(const void*, size_t) override;
     virtual KResult protocol_connect(FileDescription&, ShouldBlock) override;