/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#include <netdb.h>
#include <sys/wait.h>

#include <thrift/c_glib/transport/thrift_transport.h>
#include <thrift/c_glib/transport/thrift_socket.h>
#include <thrift/c_glib/transport/thrift_server_transport.h>
#include <thrift/c_glib/transport/thrift_server_socket.h>
#include <thrift/c_glib/transport/thrift_memory_buffer.h>

#define TEST_DATA { 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j' }

#include "../src/thrift/c_glib/transport/thrift_framed_transport.c"

static void thrift_server (const int port);
static void thrift_socket_server_open (const int port, int times);

/* test object creation and destruction */
static void
test_create_and_destroy(void)
{
  ThriftTransport *transport = NULL;
  guint r_buf_size = 0;
  guint w_buf_size = 0;

  GObject *object = NULL;
  object = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT, NULL);
  g_assert (object != NULL);
  g_object_get (G_OBJECT (object), "transport", &transport,
		"r_buf_size", &r_buf_size,
		"w_buf_size", &w_buf_size, NULL);
  g_object_unref (object);
}

static void
test_open_and_close(void)
{
  ThriftSocket *tsocket = NULL;
  ThriftTransport *transport = NULL;
  GError *err = NULL;
  pid_t pid;
  int port = 51199;
  int status;

  pid = fork ();
  g_assert ( pid >= 0 );

  if ( pid == 0 )
    {
      /* child listens */
      thrift_socket_server_open (port,1);
      exit (0);
    } else {
	/* parent connects, wait a bit for the socket to be created */
	sleep (1);
	/* create a ThriftSocket */
	tsocket = g_object_new (THRIFT_TYPE_SOCKET, "hostname", "localhost",
				"port", port, NULL);

	/* create a BufferedTransport wrapper of the Socket */
	transport = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT,
				  "transport", THRIFT_TRANSPORT (tsocket), NULL);

	/* this shouldn't work */
	g_assert (thrift_framed_transport_open (transport, NULL) == TRUE);
	g_assert (thrift_framed_transport_is_open (transport) == TRUE);
	g_assert (thrift_framed_transport_close (transport, NULL) == TRUE);
	g_object_unref (transport);
	g_object_unref (tsocket);

	/* try and underlying socket failure */
	tsocket = g_object_new (THRIFT_TYPE_SOCKET, "hostname", "localhost.broken",
				NULL);

	/* create a BufferedTransport wrapper of the Socket */
	transport = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT,
				  "transport", THRIFT_TRANSPORT (tsocket), NULL);

	g_assert (thrift_framed_transport_open (transport, &err) == FALSE);
	g_object_unref (transport);
	g_object_unref (tsocket);
	g_error_free (err);
	err = NULL;

	g_assert ( wait (&status) == pid );
	g_assert ( status == 0 );
    }
}

static void
test_read_and_write(void)
{
  int status;
  pid_t pid;
  ThriftSocket *tsocket = NULL;
  ThriftTransport *transport = NULL;
  int port = 51199;
  guchar buf[10] = TEST_DATA; /* a buffer */

  pid = fork ();
  g_assert ( pid >= 0 );

  if ( pid == 0 )
    {
      /* child listens */
      thrift_server (port);
      exit (0);
    } else {
	/* parent connects, wait a bit for the socket to be created */
	sleep (1);

	tsocket = g_object_new (THRIFT_TYPE_SOCKET, "hostname", "localhost",
				"port", port, NULL);
	transport = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT,
				  "transport", THRIFT_TRANSPORT (tsocket),
				  "w_buf_size", 4, NULL);

	g_assert (thrift_framed_transport_open (transport, NULL) == TRUE);
	g_assert (thrift_framed_transport_is_open (transport));

	/* write 10 bytes */
	thrift_framed_transport_write (transport, buf, 10, NULL);
	thrift_framed_transport_flush (transport, NULL);

	thrift_framed_transport_write (transport, buf, 1, NULL);
	thrift_framed_transport_flush (transport, NULL);

	thrift_framed_transport_write (transport, buf, 10, NULL);
	thrift_framed_transport_flush (transport, NULL);

	thrift_framed_transport_write (transport, buf, 10, NULL);
	thrift_framed_transport_flush (transport, NULL);

	thrift_framed_transport_write_end (transport, NULL);
	thrift_framed_transport_flush (transport, NULL);
	thrift_framed_transport_close (transport, NULL);

	g_object_unref (transport);
	g_object_unref (tsocket);

	g_assert ( wait (&status) == pid );
	g_assert ( status == 0 );
    }
}

/* test reading from the transport after the peer has unexpectedly
   closed the connection */
static void
test_read_after_peer_close(void)
{
  int status;
  pid_t pid;
  int port = 51199;
  GError *err = NULL;

  pid = fork ();
  g_assert (pid >= 0);

  if (pid == 0)
    {
      ThriftServerTransport *server_transport = NULL;
      ThriftTransport *client_transport = NULL;

      /* child listens */
      server_transport = g_object_new (THRIFT_TYPE_SERVER_SOCKET,
				       "port", port,
				       NULL);
      g_assert (server_transport != NULL);

      thrift_server_transport_listen (server_transport, &err);
      g_assert (err == NULL);

      /* wrap the client transport in a ThriftFramedTransport */
      client_transport = g_object_new
	  (THRIFT_TYPE_FRAMED_TRANSPORT,
	   "transport",  thrift_server_transport_accept (server_transport, &err),
	   "r_buf_size", 0,
	   NULL);
      g_assert (err == NULL);
      g_assert (client_transport != NULL);

      /* close the connection immediately after the client connects */
      thrift_transport_close (client_transport, NULL);

      g_object_unref (client_transport);
      g_object_unref (server_transport);

      exit (0);
    } else {
	ThriftSocket *tsocket = NULL;
	ThriftTransport *transport = NULL;
	guchar buf[10]; /* a buffer */

	/* parent connects, wait a bit for the socket to be created */
	sleep (1);

	tsocket = g_object_new (THRIFT_TYPE_SOCKET,
				"hostname", "localhost",
				"port",     port,
				NULL);
	transport = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT,
				  "transport",  THRIFT_TRANSPORT (tsocket),
				  "w_buf_size", 0,
				  NULL);

	g_assert (thrift_transport_open (transport, NULL) == TRUE);
	g_assert (thrift_transport_is_open (transport));

	/* attempting to read from the transport after the peer has closed
       the connection fails gracefully without generating a critical
       warning or segmentation fault */
	thrift_transport_read (transport, buf, 10, &err);
	g_assert (err != NULL);

	g_error_free (err);
	err = NULL;

	thrift_transport_read_end (transport, &err);
	g_assert (err == NULL);

	thrift_transport_close (transport, &err);
	g_assert (err == NULL);

	g_object_unref (transport);
	g_object_unref (tsocket);

	g_assert (wait (&status) == pid);
	g_assert (status == 0);
    }
}

static void
thrift_socket_server_open (const int port, int times)
{
  ThriftServerTransport *transport = NULL;
  ThriftTransport *client = NULL;
  int i;

  ThriftServerSocket *tsocket = g_object_new (THRIFT_TYPE_SERVER_SOCKET,
					      "port", port, NULL);

  transport = THRIFT_SERVER_TRANSPORT (tsocket);
  thrift_server_transport_listen (transport, NULL);
  for(i=0;i<times;i++){
      client = thrift_server_transport_accept (transport, NULL);
      g_assert (client != NULL);
      thrift_socket_close (client, NULL);
      g_object_unref (client);
  }
  g_object_unref (tsocket);
}

static void
thrift_server (const int port)
{
  int bytes = 0;
  ThriftServerTransport *transport = NULL;
  ThriftTransport *client = NULL;
  guchar buf[12]; /* a buffer */
  guchar match[10] = TEST_DATA;

  ThriftServerSocket *tsocket = g_object_new (THRIFT_TYPE_SERVER_SOCKET,
					      "port", port, NULL);

  transport = THRIFT_SERVER_TRANSPORT (tsocket);
  thrift_server_transport_listen (transport, NULL);

  /* wrap the client in a BufferedTransport */
  client = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT, "transport",
			 thrift_server_transport_accept (transport, NULL),
			 "r_buf_size", 5, NULL);
  g_assert (client != NULL);

  /* read 10 bytes */
  bytes = thrift_framed_transport_read (client, buf, 10, NULL);
  g_assert (bytes == 10); /* make sure we've read 10 bytes */
  g_assert ( memcmp (buf, match, 10) == 0 ); /* make sure what we got matches */

  bytes = thrift_framed_transport_read (client, buf, 6, NULL);
  bytes = thrift_framed_transport_read (client, buf, 5, NULL);
  bytes = thrift_framed_transport_read (client, buf, 1, NULL);

  bytes = thrift_framed_transport_read (client, buf, 12, NULL);

  thrift_framed_transport_read_end (client, NULL);
  thrift_framed_transport_close (client, NULL);
  g_object_unref (client);
  g_object_unref (tsocket);
}

/* append a framed message (4-byte big-endian size header + body) */
static void
append_frame (GByteArray *wire, const guchar *data, guint32 len)
{
  guint32 netlen = htonl (len);
  g_byte_array_append (wire, (const guchar *) &netlen, 4);
  g_byte_array_append (wire, data, len);
}

/* A read whose length crosses a frame boundary takes the read_slow() path
   while bytes are still buffered from the previous frame. That leftover used
   to be copied from the GByteArray structure itself rather than its data
   member, corrupting the result and over-reading the heap once more than a
   handful of bytes remained. Drive it through a memory buffer so the path is
   exercised without a socket peer. */
static void
test_read_across_frames (void)
{
  guchar f1[100];
  guchar f2[100];
  guchar buf[100];
  gint32 got;
  guint i;

  for (i = 0; i < sizeof (f1); i++)
    {
      f1[i] = (guchar) (0x10 + i);
      f2[i] = (guchar) (0xc0 + i);
    }

  GByteArray *wire = g_byte_array_new ();
  append_frame (wire, f1, sizeof (f1));
  append_frame (wire, f2, sizeof (f2));

  ThriftMemoryBuffer *membuf = g_object_new (THRIFT_TYPE_MEMORY_BUFFER,
                                             "buf", wire,
                                             "buf_size", (guint32) 0,
                                             NULL);
  ThriftTransport *transport = g_object_new (THRIFT_TYPE_FRAMED_TRANSPORT,
                                             "transport",
                                             THRIFT_TRANSPORT (membuf),
                                             NULL);

  /* consume part of the first frame so 96 bytes stay buffered */
  got = thrift_transport_read (transport, buf, 4, NULL);
  g_assert (got == 4);
  g_assert (memcmp (buf, f1, 4) == 0);

  /* this read spans into the second frame and must return the real buffered
     bytes, not the bytes of the GByteArray structure */
  got = thrift_transport_read (transport, buf, 100, NULL);
  g_assert (got == 100);
  g_assert (memcmp (buf, f1 + 4, 96) == 0);
  g_assert (memcmp (buf + 96, f2, 4) == 0);

  thrift_transport_read_end (transport, NULL);
  g_object_unref (transport);
  g_object_unref (membuf);
}

int
main(int argc, char *argv[])
{
#if (!GLIB_CHECK_VERSION (2, 36, 0))
  g_type_init();
#endif

  g_test_init (&argc, &argv, NULL);

  g_test_add_func ("/testframedtransport/CreateAndDestroy", test_create_and_destroy);
  g_test_add_func ("/testframedtransport/OpenAndClose", test_open_and_close);
  g_test_add_func ("/testframedtransport/ReadAndWrite", test_read_and_write);
  g_test_add_func ("/testframedtransport/ReadAfterPeerClose", test_read_after_peer_close);
  g_test_add_func ("/testframedtransport/ReadAcrossFrames", test_read_across_frames);

  return g_test_run ();
}
