#!/usr/bin/perl -w
# Copyright © 2018 Jamie Zawinski <jwz@dnalounge.com>
#
# Permission to use, copy, modify, distribute, and sell this software and its
# documentation for any purpose is hereby granted without fee, provided that
# the above copyright notice appear in all copies and that both that
# copyright notice and this permission notice appear in supporting
# documentation.  No representations are made about the suitability of this
# software for any purpose.  It is provided "as is" without express or 
# implied warranty.
#
# This is an extremely simple WebSocket server for DNA Lounge internal apps.
# 
#   - Connections are denied unless the https set-up presents a valid
#     DNA_AUTH cookie (which means, JavaScript can only connect to this
#     server if the invoking browser has already logged in).
#
#   - Except that's not required if the client IP is the server IP.
#
#   - Any messages you send are sent to every other participant who connected
#     with the same path part. (Messages are not echoed back to yourself).
#
#   - You say "ping" I say "pong".
#
#   - For SSL, use: --ssl --certs privkey.pem fullchain.pem
#     Note that browsers don't like self-signed certs in websockets.
#
#   - To restrict to same-origin, use: --origin yourdomain.com
#     This means that connections will be rejected when initiated from
#     JavaScript that was not served from your domain, to avoid cross-site
#     hijacking and exfiltration of your authentication cookies.
#
# Created:  6-Sep-2018.

require 5;
use diagnostics;
use strict;

use POSIX;
use Socket;
use IO::Socket::SSL;
use Net::WebSocket::Server;
use Sys::Hostname;

# Not yours
BEGIN { push @INC, "/var/www/dnalounge/utils/"; }
use dna_auth;

use open ":encoding(utf8)";

my $progname = $0; $progname =~ s@.*/@@g;
my ($version) = ('$Revision: 1.10 $' =~ m/\s(\d[.\d]+)\s/s);

my $verbose = 0;
my $debug_p = 0;

my $server = undef;

# If we don't call shutdown() upon signals and other abnormal exits, then
# the kernel keeps listening on our port number for up to 30 seconds, and
# we can't re-launch until that times out.  How very.
#
# But I guess sometimes that happens anyway? WTF.
#
END {
  LOG ("END") if ($verbose);
  $server->shutdown() if $server;
}

sub signal_cleanup($) {
  my ($s) = @_;
  LOG ("SIG$s");
  $server->shutdown() if $server;
  $server = undef;
  exit (2);  # This causes END{} to run.
}

$SIG{HUP}  = \&signal_cleanup;
$SIG{INT}  = \&signal_cleanup;
$SIG{QUIT} = \&signal_cleanup;
$SIG{ABRT} = \&signal_cleanup;
$SIG{KILL} = \&signal_cleanup;
$SIG{TERM} = \&signal_cleanup;


sub url_unquote($) {
  my ($url) = @_;
  $url =~ s/[+]/ /g;
  $url =~ s/%([a-z0-9]{2})/chr(hex($1))/ige;
  return $url;
}


# Per-channel source IP whitelists.
my %whitelists = ( '/jwztv' => '^(50\.0\.18\.32|75\.101\.62\.19)$'
                 );


sub launch_server($$;$$$) {
  my ($sslp, $port, $key, $crt, $origin) = @_;

  LOG ("listening on port $port") if ($verbose);

  $IO::Socket::SSL::DEBUG = $verbose-1 if ($verbose);
  $SSL_ERROR = '';

  error ("origin should be a domain name: $origin")
    if (defined($origin) && !($origin =~ m/^[-_a-z\d]+(\.[-_a-z\d]+)+$/s));

  my $S = ($sslp
           ? IO::Socket::SSL->new (Listen         => 5,
                                   LocalPort      => $port,
                                   SSL_cert_file  => $crt,
                                   SSL_key_file   => $key,
                                   ReuseAddr      => 1)
           : IO::Socket::INET->new (Listen        => 5,
                                   LocalPort      => $port,
                                   ReuseAddr      => 1));
  error ("socket failed: $! $SSL_ERROR") unless $S;

  # Try to avoid "Address already in use" when restarting.
  # Since END calls shutdown, this shouldn't be necessary.
  # Also, this doesn't help. The fucking kernel is still marking our
  # port as "in use" for 30+ seconds after we have cleanly exited,
  # kernel 2.6.32 Centos 6.2. How very.
  #
  # Need to do this in ->new instead, maybe?
  #
  # setsockopt ($S, SOL_SOCKET, SO_REUSEADDR, 1) ||
  #    LOG ("SO_REUSEADDR failed: $!");

  my $server_ip = inet_ntoa ((gethostbyname (hostname()))[4]);

  $server = Net::WebSocket::Server->new (listen => $S,
    on_connect => sub {
      my ($serv, $conn) = @_;

      $conn->on (
        handshake => sub {
          my ($conn, $handshake) = @_;

          my $client_ip = $conn->ip();
          my $client_port = $conn->port();

          # Save the path part of the URL into $conn so that we can tell
          # which connections are a part of the same "room".
          my $res = $handshake->req->resource_name || '';
          $conn->{resource_name} = $res;

          # Save the IP and port for later diagnostics, since disconnect()
          # doesn't receive a connection object that still has an IP.
          my $cid = "$client_ip:$client_port$res";
          $conn->{cid} = $cid;

          # Make sure the JS code is hosted from our domain.
          if ($origin) {
            my $ourl = $handshake->res->origin || '';
            my ($ohost) = ($ourl =~ m@^[^:/]+://([^:/]+)@s);
            if ($ohost && $ohost =~ m/\b\Q$origin\E$/si) {
              LOG ("$cid: origin matches: $origin, $ourl") if ($verbose);
            } else {
              LOG ("$cid: origin mismatch: $origin, $ourl") if ($verbose);
              $conn->disconnect();
              return;
            }
          }

          my $auth_required_p = !$debug_p;
          my $re = $whitelists{$res};

          if ($client_ip eq $server_ip ||
              $client_ip eq '127.0.0.1' ||
              $client_ip eq '70.36.236.109' ||  # staging.dnalounge.com
              $re && $client_ip =~ m/$re/s) {
            LOG ("$cid: localhost") if ($verbose);
            $auth_required_p = 0;
          }

          my $cookie = undef;
          if ($handshake->req &&
              $handshake->req->cookies &&
              ref($handshake->req->cookies) eq 'Protocol::WebSocket::Cookie' &&
              ref($handshake->req->cookies->pairs)) {
            foreach my $pair (@{$handshake->req->cookies->pairs}) {
              if ($pair->[0] eq 'DNA_AUTH') {
                $cookie = url_unquote ($pair->[1]);
                last;
              }
            }
          }

          if (!$cookie && $auth_required_p) {
            LOG ("$cid: handshake: no DNA cookie");
            $conn->disconnect();
            return;
          }

          if ($cookie &&
              !dna_auth::dna_auth_validate_login_cookie ($cookie)) {
              LOG ("$cid: handshake: invalid DNA cookie" .
                   ($debug_p ? ', DEBUG' : ''));
            if ($auth_required_p) {
              $conn->disconnect();
              return;
            }
          }

          LOG ("$cid: handshake: authorized") if ($verbose);
        },

        ready => sub {
          my ($conn) = @_;
          LOG ($conn->{cid} . ": ready") if ($verbose);
        },

        utf8 => sub {
          my ($conn, $msg) = @_;
          my $res = $conn->{resource_name};

          LOG ($conn->{cid} . ": utf8: $msg") if ($verbose);

          if ($msg =~ m/^ping\s*$/s) {
            $conn->send_utf8 ("pong\n");
            return;
          }

          foreach my $conn2 ($conn->server->connections) {
            if ($conn ne $conn2 &&
                $res eq $conn2->{resource_name}) {
              $conn2->send_utf8 ($msg);
            }
          }
        },

        binary => sub {
          my ($conn, $msg) = @_;
          my $res = $conn->{resource_name};
          LOG ($conn->{cid} . ": binary: " . length($msg) . " bytes")
            if ($verbose);
          foreach my $conn2 ($conn->server->connections) {
            if ($conn ne $conn2 &&
                $res eq $conn2->{resource_name}) {
              $conn2->send_binary ($msg);
            }
          }
        },

        pong => sub {
          my ($conn) = @_;
          LOG ($conn->{cid} . ": pong") if ($verbose);
        },

        disconnect => sub {
          my ($conn) = @_;
          LOG ($conn->{cid} . ": disconnect") if ($verbose);
        },
      ),
    },

    on_tick => sub {
      my ($serv) = @_;
      LOG ("tick") if ($verbose);
    },

    on_shutdown => sub {
      my ($serv) = @_;
      LOG ("shutdown") if ($verbose);

      # Try closing each client explicitly.  But only once.
      # Maybe this fixes the "Address already in use" problem?
      #
      if (! $serv->{shutting_down_p}) {
        $serv->{shutting_down_p} = 1;
        foreach my $conn ($serv->connections) {
          LOG ($conn->{cid} . ": shutdown disconnect") if ($verbose);
          $conn->disconnect (1001);
          $conn->{listen}->close() if ($conn->{listen});
        }
      }
    },
  );

  eval { $server->start(); };
  LOG ("error: $@") if ($@);
  exit (1);  # run END and call shutdown()
}


sub LOG($) {
  my ($s) = @_;
  print STDERR "$progname: " . strftime("%H:%M:%S", localtime) . ": $s\n";
  system ('logger', '-t', $progname, $s);
}

sub error($) {
  my ($err) = @_;
  LOG ("ERROR: $err");
  exit 1;
}

sub usage() {
  print STDERR "usage: $progname [--verbose] [--ssl] [--port P]" .
    " [--origin DOMAIN]" .
    " [--certs PRIVKEY FULLCHAIN]\n";
  exit 1;
}

sub main() {
  my $port = undef;
  my ($sslp, $crt, $key, $origin);
  while (@ARGV) {
    $_ = shift @ARGV;
    if (m/^--?verbose$/) { $verbose++; }
    elsif (m/^-v+$/) { $verbose += length($_)-1; }
    elsif (m/^--?debug$/) { $debug_p++; }
    elsif (m/^--?port$/) { $port = 0 + shift @ARGV; }
    elsif (m/^--?ssl$/) { $sslp = 1; }
    elsif (m/^--?certs?$/) { $key = shift @ARGV; $crt = shift @ARGV; }
    elsif (m/^--?origin$/) { $origin = shift @ARGV; }
    elsif (m/^-./) { usage; }
    else { usage; }
  }

  usage unless ($port && $port > 80);

  launch_server ($sslp, $port, $key, $crt, $origin);
}

main();
exit 0;