/*

  sshserverprobe.c

  Authors:
        Timo J. Rinne <tri@ssh.com>

  Copyright (C) 1999-2000 SSH Communications Security Corp, Helsinki, Finland
  All rights reserved.

*/

#include "ssh2includes.h"
#include "ssheloop.h"
#include "sshtimeouts.h"
#include "sshtcp.h"
#include "sshudp.h"
#include "sshbuffer.h"
#include "sshserver.h"
#include "sshencode.h"
#include "sshserverprobe.h"

#define SSH_DEBUG_MODULE "SshServerProbe" 

#ifdef __SUNPRO_C
#pragma error_messages (off,E_STATEMENT_NOT_REACHED)
#endif /* __SUNPRO_C */

typedef struct SshServerProbeRec *SshServerProbe;

struct SshServerProbeRec {
  SshOperationHandle handle;
  SshUdpListener udp;
  SshServerProbeCB callback;
  Boolean in_udp_callback;
  Boolean aborted;
  void *context;
};

/* Normal timeout */
void ssh_server_probe_timeout(void *context)
{
  SshServerProbe probe = (SshServerProbe)context;

  ssh_operation_unregister(probe->handle);
  ssh_udp_destroy_listener(probe->udp);
  if (probe->callback)
    probe->callback(NULL, NULL, NULL, probe->context);
  ssh_xfree(probe);
  return;
}

/* If this function is called from the udp_callback, the context is
   not freed.  Instead aborted flag is risen and udp callback frees
   the context and returns control to the user. */
void ssh_server_probe_abort(void *context)
{
  SshServerProbe probe = (SshServerProbe)context;

  ssh_cancel_timeouts(ssh_server_probe_timeout, (void *)probe);
  ssh_udp_destroy_listener(probe->udp);
  if (probe->in_udp_callback)
    ssh_xfree(probe);
  else
    probe->aborted = TRUE;
  return;
}

/* UDP callback. */
void ssh_server_probe_udp_callback(SshUdpListener listener, void *context)
{
  SshServerProbe probe = (SshServerProbe)context;
  SshUdpError err;
  char remote_address[64];
  char remote_port[32];
  unsigned char packet[2048];
  size_t packet_len, consumed, len;
  char *remote_command, *remote_version, *remote_server_port;

  while (1)
    {
      err = ssh_udp_read(listener,
                         remote_address, sizeof (remote_address),
                         remote_port, sizeof (remote_port),
                         packet, sizeof (packet),
                         &packet_len);
      remote_command = NULL;
      remote_version = NULL;
      remote_server_port = NULL;
      switch (err)
        {
        case SSH_UDP_OK:
          SSH_DEBUG(5, ("addr=%s port=%s len=%d",
                        remote_address,
                        remote_port,
                        packet_len));
          len = ssh_decode_array(packet, packet_len,
                                 SSH_FORMAT_UINT32_STR, 
                                 &remote_command, NULL,
                                 SSH_FORMAT_END);
          consumed = len;
          if (consumed == 0)
            {
              SSH_DEBUG(5, ("malformed reply"));
              break;
            }
          if (strcmp(remote_command, "server-reply") == 0)
            {
              len = ssh_decode_array(packet + consumed, 
                                     packet_len - consumed,
                                     SSH_FORMAT_UINT32_STR, 
                                     &remote_version, NULL,
                                     SSH_FORMAT_UINT32_STR, 
                                     &remote_server_port, NULL,
                                     SSH_FORMAT_END);
              if (len == 0)
                {
                  SSH_DEBUG(5, ("malformed server-reply"));
                  break;
                }
              else
                {
                  consumed += len;
                }
              /* User may call ssh_operation_abort from this callback.
                 This flag is for the abort callback so that it doesn't
                 free the context. */
              probe->in_udp_callback = TRUE;
              if (probe->callback)
                probe->callback(remote_address, 
                                remote_server_port, 
                                remote_version, 
                                probe->context);
              probe->in_udp_callback = FALSE;
              if (probe->aborted)
                {
                  /* User called ssh_operation_abort from callback.
                     Context is otherwise invalid, we just free the
                     struct here and return.  No more callbacks are
                     allowed. */
                  ssh_xfree(probe);
                  return;
                }
            }
          else
            {
              SSH_DEBUG(5, ("unknown reply"));
            }
          break;

        case SSH_UDP_HOST_UNREACHABLE:
        case SSH_UDP_PORT_UNREACHABLE:
          break;

        case SSH_UDP_NO_DATA:
          return;
        }
      ssh_xfree(remote_command);
      ssh_xfree(remote_version);
      ssh_xfree(remote_server_port);
    }
  SSH_NOTREACHED;
}

/* Send a server probe datragram to the broadcast address defined by
   the caller.  Call probe_cb with address, port and version string of
   each server that replied.  If udp datagram cannot be sent at all,
   callback may be called immediately with NULL address, port and
   version.  In this case returned handle is null and no further
   callbacks are called.  In normal operation, valid operation handle
   is returned and operation can be aborted with ssh_operation_abort.
   When timeout is exceeded, the callback is called with arguments
   (NULL, NULL, NULL, context).  After this operation handle is
   invalid. */
SshOperationHandle ssh_server_probe(const char *broadcast_address,
                                    long timeout_sec,
                                    long timeout_usec,
                                    SshServerProbeCB probe_cb,
                                    void *context)
{
  SshServerProbe probe;
  SshBufferStruct buffer[1];

  probe = ssh_xcalloc(1, sizeof (*probe));
  probe->handle = ssh_operation_register(ssh_server_probe_abort, probe);
  probe->callback = probe_cb;
  probe->context = context;
  probe->udp = ssh_udp_make_listener(NULL,   
                                     NULL,
                                     NULL,
                                     NULL,
                                     ssh_server_probe_udp_callback,
                                     (void *)probe);
  if (probe->udp == NULL)
    {
      if (probe->callback)
        probe->callback(NULL, NULL, NULL, probe->context);
      ssh_operation_unregister(probe->handle);
      ssh_xfree(probe);
      return NULL;
    }
  ssh_udp_set_broadcasting(probe->udp, TRUE);
  ssh_buffer_init(buffer);
  ssh_encode_buffer(buffer,
                    SSH_FORMAT_UINT32_STR, 
                    "server-query", 
                    strlen ("server-query"),
                    SSH_FORMAT_UINT32_STR, 
                    SSH2_VERSION_STRING, 
                    strlen(SSH2_VERSION_STRING),
                    SSH_FORMAT_END);
  ssh_udp_send(probe->udp, 
               broadcast_address, SSH_DEFAULT_PORT, 
               ssh_buffer_ptr(buffer),
               ssh_buffer_len(buffer));
  ssh_buffer_uninit(buffer);
  ssh_register_timeout(timeout_sec, timeout_usec, 
                       ssh_server_probe_timeout, (void *)probe);
  return probe->handle;
}

/* eof (sshserverprobe.c) */

