Add encrypted LAN mouse and keyboard sharing with four-edge handoff

src/lan_channel.cpp

@@ -0,0 +1,185 @@
+#include <winsock2.h>
+#include <ws2tcpip.h>
+#include "lan_channel.h"
+#include <bcrypt.h>
+#include <limits>
+
+namespace ks::lan
+{
+namespace
+{
+void checked(NTSTATUS s){if(s<0) throw std::runtime_error("Input channel authentication failed.");}
+void put32(Bytes &b,uint32_t v){for(int n=24;n>=0;n-=8)b.push_back(static_cast<unsigned char>(v>>n));}
+uint32_t get32(const unsigned char *p){return (uint32_t(p[0])<<24)|(uint32_t(p[1])<<16)|(uint32_t(p[2])<<8)|p[3];}
+Bytes keyBytes(const std::string &key)
+{
+    if(key.size()!=64||key.find_first_not_of("0123456789abcdef")!=key.npos) throw std::runtime_error("Invalid input key.");
+    Bytes bytes;for(size_t i=0;i<key.size();i+=2)bytes.push_back(static_cast<unsigned char>(std::stoul(key.substr(i,2),nullptr,16)));return bytes;
+}
+bool nonblocking(SOCKET s){u_long value=1;BOOL yes=TRUE;return !ioctlsocket(s,FIONBIO,&value)&&!setsockopt(s,IPPROTO_TCP,TCP_NODELAY,(char*)&yes,sizeof(yes));}
+void closeSocket(SOCKET &s){if(s!=INVALID_SOCKET){closesocket(s);s=INVALID_SOCKET;}}
+}
+struct InputCipher::Impl
+{
+    BCRYPT_ALG_HANDLE alg{};BCRYPT_KEY_HANDLE key{};uint64_t sequence{};
+    explicit Impl(const std::string &secret)
+    {
+        auto bytes=keyBytes(secret);
+        try {
+            checked(BCryptOpenAlgorithmProvider(&alg,BCRYPT_AES_ALGORITHM,nullptr,0));
+            checked(BCryptSetProperty(alg,BCRYPT_CHAINING_MODE,(PUCHAR)BCRYPT_CHAIN_MODE_GCM,sizeof(BCRYPT_CHAIN_MODE_GCM),0));
+            checked(BCryptGenerateSymmetricKey(alg,&key,nullptr,0,bytes.data(),ULONG(bytes.size()),0));
+        } catch(...) {if(key)BCryptDestroyKey(key);if(alg)BCryptCloseAlgorithmProvider(alg,0);throw;}
+        SecureZeroMemory(bytes.data(),bytes.size());
+    }
+    ~Impl(){if(key)BCryptDestroyKey(key);if(alg)BCryptCloseAlgorithmProvider(alg,0);}
+    Bytes crypt(Bytes payload,Bytes &tag,bool encrypt)
+    {
+        if(sequence==std::numeric_limits<uint64_t>::max())throw std::runtime_error("Input session exhausted.");
+        unsigned char nonce[12]{};for(int i=0;i<8;++i)nonce[11-i]=static_cast<unsigned char>(sequence>>(i*8));
+        BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO info;BCRYPT_INIT_AUTH_MODE_INFO(info);
+        info.pbNonce=nonce;info.cbNonce=sizeof(nonce);info.pbTag=tag.data();info.cbTag=ULONG(tag.size());
+        // Protocol/direction are bound by key derivation; sequence and payload length by GCM.
+        Bytes result(payload.size());ULONG written{};
+        auto fn=encrypt?BCryptEncrypt:BCryptDecrypt;
+        checked(fn(key,payload.data(),ULONG(payload.size()),&info,nullptr,0,result.data(),ULONG(result.size()),&written,0));
+        if(written!=payload.size())throw std::runtime_error("Invalid input frame size.");
+        ++sequence;return result;
+    }
+};
+InputCipher::InputCipher(const std::string &key):impl(std::make_unique<Impl>(key)){}
+InputCipher::~InputCipher()=default;
+Bytes InputCipher::encrypt(const InputFrame &f)
+{
+    Bytes plain;put32(plain,uint32_t(f.type));put32(plain,f.token);put32(plain,uint32_t(f.a));put32(plain,uint32_t(f.b));put32(plain,uint32_t(f.c));
+    Bytes tag(16);auto result=impl->crypt(std::move(plain),tag,true);result.insert(result.end(),tag.begin(),tag.end());return result;
+}
+InputFrame InputCipher::decrypt(const Bytes &packet)
+{
+    if(packet.size()!=36)throw std::runtime_error("Invalid encrypted input frame.");
+    Bytes tag(packet.begin()+20,packet.end());auto plain=impl->crypt(Bytes(packet.begin(),packet.begin()+20),tag,false);
+    auto type=get32(plain.data());if(type>uint32_t(InputType::Wheel))throw std::runtime_error("Unsupported input frame.");
+    return {InputType(type),get32(plain.data()+4),int32_t(get32(plain.data()+8)),int32_t(get32(plain.data()+12)),int32_t(get32(plain.data()+16))};
+}
+std::string inputSessionKey(const std::string &secret,const std::string &from,const std::string &to,const std::string &a,const std::string &b,int edge)
+{
+    if(!identifier(from)||!identifier(to)||from==to||!identifier(a)||!identifier(b))throw std::runtime_error("Invalid session identity.");
+    return mac(secret,"KSwitcher input v1|"+from+"|"+to+"|"+a+"|"+b+"|"+std::to_string(edge));
+}
+struct InputChannel::Impl
+{
+    std::string local,peer,secret,nonce,state="Waiting for the paired computer...";
+    int edge;uint16_t port;SOCKET listener=INVALID_SOCKET,socket=INVALID_SOCKET;
+    bool connecting{},authenticated{},helloRead{};
+    uint64_t changed{},lastRead{},lastPing{},retry{};
+    Bytes incoming,outgoing;
+    std::unique_ptr<InputCipher> tx,rx;
+    Impl(std::string a,std::string b,std::string key,int side,uint16_t p):local(std::move(a)),peer(std::move(b)),secret(std::move(key)),edge(side),port(p){}
+    ~Impl(){closeSocket(socket);closeSocket(listener);}
+    void reset(uint64_t now,const std::string &reason)
+    {
+        closeSocket(socket);connecting=authenticated=helloRead=false;incoming.clear();outgoing.clear();tx.reset();rx.reset();++changed;retry=now+1500;state=reason;
+    }
+    void begin(uint64_t now)
+    {
+        nonce=randomHex(16);lastRead=lastPing=now;
+        auto hello=Json{{"app","KSwitcher-input"},{"version",1},{"id",local},{"to",peer},{"nonce",nonce},{"edge",edge}}.dump();
+        put32(outgoing,uint32_t(hello.size()));outgoing.insert(outgoing.end(),hello.begin(),hello.end());
+        state="Authenticating paired KSwitcher...";
+    }
+    void poll(uint32_t address,const std::vector<InputFrame> &frames,std::vector<InputFrame> &received,uint64_t now)
+    {
+        try
+        {
+            if(socket==INVALID_SOCKET&&now<retry)return;
+            if(listener==INVALID_SOCKET && local>peer)
+            {
+                listener=::socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);BOOL yes=TRUE;
+                if(listener==INVALID_SOCKET||setsockopt(listener,SOL_SOCKET,SO_EXCLUSIVEADDRUSE,(char*)&yes,sizeof(yes))||!nonblocking(listener)){closeSocket(listener);throw std::runtime_error("Cannot open input listener.");}
+                sockaddr_in a{};a.sin_family=AF_INET;a.sin_port=htons(port);
+                if(bind(listener,(sockaddr*)&a,sizeof(a))||listen(listener,1)) {closeSocket(listener);throw std::runtime_error("TCP input port is unavailable.");}
+            }
+            if(socket==INVALID_SOCKET && address && now>=retry)
+            {
+                if(local<peer)
+                {
+                    socket=::socket(AF_INET,SOCK_STREAM,IPPROTO_TCP);if(socket==INVALID_SOCKET||!nonblocking(socket))throw std::runtime_error("Cannot open input connection.");
+                    sockaddr_in a{};a.sin_family=AF_INET;a.sin_port=htons(port);a.sin_addr.s_addr=address;
+                    int result=connect(socket,(sockaddr*)&a,sizeof(a));
+                    if(result && WSAGetLastError()!=WSAEWOULDBLOCK)throw std::runtime_error("Paired computer is not accepting input connections.");
+                    connecting=result!=0;begin(now);
+                }
+                else
+                {
+                    sockaddr_in a{};int length=sizeof(a);socket=accept(listener,(sockaddr*)&a,&length);
+                    if(socket!=INVALID_SOCKET)
+                    {
+                        if(a.sin_addr.s_addr!=address){closeSocket(socket);return;}
+                        if(!nonblocking(socket))throw std::runtime_error("Cannot configure input connection.");
+                        begin(now);
+                    }
+                }
+            }
+            if(socket==INVALID_SOCKET)return;
+            if(now-lastRead>2000)throw std::runtime_error("Input connection timed out; control returned locally.");
+            if(connecting)
+            {
+                fd_set writes,errors;FD_ZERO(&writes);FD_ZERO(&errors);FD_SET(socket,&writes);FD_SET(socket,&errors);timeval wait{};
+                if(select(0,nullptr,&writes,&errors,&wait)<=0)return;
+                int error=0;int size=sizeof(error);if(getsockopt(socket,SOL_SOCKET,SO_ERROR,(char*)&error,&size)||error)throw std::runtime_error("Cannot connect to paired KSwitcher.");
+                connecting=false;
+            }
+            if(authenticated)
+            {
+                for(auto &f:frames){auto bytes=tx->encrypt(f);outgoing.insert(outgoing.end(),bytes.begin(),bytes.end());}
+                if(now-lastPing>=400){auto bytes=tx->encrypt({InputType::Ping});outgoing.insert(outgoing.end(),bytes.begin(),bytes.end());lastPing=now;}
+            }
+            if(outgoing.size()>32768)throw std::runtime_error("Input queue overflow; control returned locally.");
+            if(!outgoing.empty())
+            {
+                int count=::send(socket,(char*)outgoing.data(),int(outgoing.size()),0);
+                if(count>0)outgoing.erase(outgoing.begin(),outgoing.begin()+count);
+                else if(count==0||WSAGetLastError()!=WSAEWOULDBLOCK)throw std::runtime_error("Input send failed.");
+            }
+            unsigned char data[4096];
+            for(int batch=0;batch<8;++batch)
+            {
+                int count=recv(socket,(char*)data,sizeof(data),0);
+                if(count==0)throw std::runtime_error("Other KSwitcher disconnected.");
+                if(count<0){if(WSAGetLastError()!=WSAEWOULDBLOCK)throw std::runtime_error("Input receive failed.");break;}
+                incoming.insert(incoming.end(),data,data+count);
+            }
+            if(incoming.size()>32768)throw std::runtime_error("Input receive queue overflow.");
+            if(!helloRead && incoming.size()>=4)
+            {
+                size_t length=get32(incoming.data());if(length>512||length<32)throw std::runtime_error("Invalid input handshake.");
+                if(incoming.size()<length+4)return;
+                auto hello=Json::parse(incoming.begin()+4,incoming.begin()+4+length);
+                if(hello.value("app","")!="KSwitcher-input"||hello.value("version",0)!=1||hello.value("id","")!=peer||hello.value("to","")!=local)
+                    throw std::runtime_error("Unexpected input peer.");
+                if(hello.value("edge",-1)!=(edge^1))throw std::runtime_error("Screen positions must be opposite on the two computers.");
+                auto remoteNonce=hello.value("nonce","");
+                tx=std::make_unique<InputCipher>(inputSessionKey(secret,local,peer,nonce,remoteNonce,edge));
+                rx=std::make_unique<InputCipher>(inputSessionKey(secret,peer,local,remoteNonce,nonce,edge^1));
+                incoming.erase(incoming.begin(),incoming.begin()+4+length);helloRead=true;
+                auto bytes=tx->encrypt({InputType::Ready});outgoing.insert(outgoing.end(),bytes.begin(),bytes.end());
+            }
+            while(helloRead && incoming.size()>=36)
+            {
+                auto f=rx->decrypt(Bytes(incoming.begin(),incoming.begin()+36));incoming.erase(incoming.begin(),incoming.begin()+36);lastRead=now;
+                if(!authenticated){if(f.type!=InputType::Ready)throw std::runtime_error("Input received before authentication.");authenticated=true;state="Input connected. Move through the configured screen edge.";}
+                else if(f.type==InputType::Ready)throw std::runtime_error("Repeated input handshake.");
+                else if(f.type!=InputType::Ping)received.push_back(f);
+            }
+        }
+        catch(const std::exception &e){received.clear();reset(now,e.what());}
+    }
+};
+InputChannel::InputChannel(std::string a,std::string b,std::string key,int edge,uint16_t port):impl(std::make_unique<Impl>(std::move(a),std::move(b),std::move(key),edge,port)){}
+InputChannel::~InputChannel()=default;
+void InputChannel::poll(uint32_t ip,const std::vector<InputFrame>&out,std::vector<InputFrame>&in,uint64_t now){impl->poll(ip,out,in,now);}
+bool InputChannel::ready()const{return impl->authenticated;}
+uint64_t InputChannel::generation()const{return impl->changed;}
+std::string InputChannel::status()const{return impl->state;}
+void InputChannel::disconnect(){impl->reset(GetTickCount64(),"Input disconnected; local control restored.");}
+}