/* -*- Mode: C; tab-width: 8; indent-tabs-mode: t; c-basic-offset: 8 -*- */

/*

 * Copyright 2013 Red Hat, Inc.

 */

#include "test-utils.h"

static void

force_io_streams_init (void)

{

	SoupServer *server;

	SoupSession *session;

	SoupURI *base_uri;

	SoupMessage *msg;

	/* Poke libsoup enough to cause SoupBodyInputStream and

	 * SoupBodyOutputStream to get defined, so we can find them

	 * via g_type_from_name() later.

	 */

	server = soup_test_server_new (TRUE);

	base_uri = soup_test_server_get_uri (server, "http", NULL);

	session = soup_test_session_new (SOUP_TYPE_SESSION, NULL);

	msg = soup_message_new_from_uri ("POST", base_uri);

	soup_session_send_message (session, msg);

	g_object_unref (msg);

	soup_test_session_abort_unref (session);

	soup_uri_free (base_uri);

	soup_test_server_quit_unref (server);

}

typedef struct {

	GFilterInputStream grandparent;

	gpointer *soup_filter_input_stream_private;

	gboolean is_readable;

} SlowInputStream;

typedef struct {

	GFilterInputStreamClass grandparent;

} SlowInputStreamClass;

GType slow_input_stream_get_type (void);

static void slow_pollable_input_stream_init (GPollableInputStreamInterface *pollable_interface,

					     gpointer interface_data);

G_DEFINE_TYPE_WITH_CODE (SlowInputStream, slow_input_stream,

			 g_type_from_name ("SoupFilterInputStream"),

			 G_IMPLEMENT_INTERFACE (G_TYPE_POLLABLE_INPUT_STREAM, slow_pollable_input_stream_init);

			 )

static void

slow_input_stream_init (SlowInputStream *sis)

{

}

static gssize

slow_input_stream_read (GInputStream  *stream,

			void          *buffer,

			gsize          count,

			GCancellable  *cancellable,

			GError       **error)

{

	return g_input_stream_read (G_FILTER_INPUT_STREAM (stream)->base_stream,

				    buffer, 1, cancellable, error);

}

static void

slow_input_stream_class_init (SlowInputStreamClass *sisclass)

{

	GInputStreamClass *input_stream_class = G_INPUT_STREAM_CLASS (sisclass);

	input_stream_class->read_fn = slow_input_stream_read;

}

static gboolean

slow_input_stream_is_readable (GPollableInputStream *stream)

{

	return ((SlowInputStream *)stream)->is_readable;

}

static gssize

slow_input_stream_read_nonblocking (GPollableInputStream  *stream,

				    void                  *buffer,

				    gsize                  count,

				    GError               **error)

{

	if (((SlowInputStream *)stream)->is_readable) {

		((SlowInputStream *)stream)->is_readable = FALSE;

		return slow_input_stream_read (G_INPUT_STREAM (stream), buffer, count,

					       NULL, error);

	} else {

		g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_WOULD_BLOCK,

				     "would block");

		return -1;

	}

}

static GSource *

slow_input_stream_create_source (GPollableInputStream *stream,

				 GCancellable *cancellable)

{

	GSource *base_source, *pollable_source;

	((SlowInputStream *)stream)->is_readable = TRUE;

	base_source = g_timeout_source_new (0);

	g_source_set_dummy_callback (base_source);

	pollable_source = g_pollable_source_new (G_OBJECT (stream));

	g_source_add_child_source (pollable_source, base_source);

	g_source_unref (base_source);

	return pollable_source;

}

static void

slow_pollable_input_stream_init (GPollableInputStreamInterface *pollable_interface,

				 gpointer interface_data)

{

	pollable_interface->is_readable = slow_input_stream_is_readable;

	pollable_interface->read_nonblocking = slow_input_stream_read_nonblocking;

	pollable_interface->create_source = slow_input_stream_create_source;

}

typedef struct {

	GFilterOutputStream parent;

	gboolean is_writable;

} SlowOutputStream;

typedef struct {

	GFilterOutputStreamClass parent;

} SlowOutputStreamClass;

GType slow_output_stream_get_type (void);

static void slow_pollable_output_stream_init (GPollableOutputStreamInterface *pollable_interface,

					      gpointer interface_data);

G_DEFINE_TYPE_WITH_CODE (SlowOutputStream, slow_output_stream,

			 g_type_from_name ("GFilterOutputStream"),

			 G_IMPLEMENT_INTERFACE (G_TYPE_POLLABLE_OUTPUT_STREAM, slow_pollable_output_stream_init);

			 )

static void

slow_output_stream_init (SlowOutputStream *sis)

{

}

static gssize

slow_output_stream_write (GOutputStream  *stream,

			  const void     *buffer,

			  gsize           count,

			  GCancellable   *cancellable,

			  GError        **error)

{

	return g_output_stream_write (G_FILTER_OUTPUT_STREAM (stream)->base_stream,

				      buffer, 1, cancellable, error);

}

static void

slow_output_stream_class_init (SlowOutputStreamClass *sisclass)

{

	GOutputStreamClass *output_stream_class = G_OUTPUT_STREAM_CLASS (sisclass);

	output_stream_class->write_fn = slow_output_stream_write;

}

static gboolean

slow_output_stream_is_writable (GPollableOutputStream *stream)

{

	return ((SlowOutputStream *)stream)->is_writable;

}

static gssize

slow_output_stream_write_nonblocking (GPollableOutputStream  *stream,

				      const void             *buffer,

				      gsize                   count,

				      GError                **error)

{

	if (((SlowOutputStream *)stream)->is_writable) {

		((SlowOutputStream *)stream)->is_writable = FALSE;

		return slow_output_stream_write (G_OUTPUT_STREAM (stream), buffer, count,

						 NULL, error);

	} else {

		g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_WOULD_BLOCK,

				     "would block");

		return -1;

	}

}

static GSource *

slow_output_stream_create_source (GPollableOutputStream *stream,

				  GCancellable *cancellable)

{

	GSource *base_source, *pollable_source;

	((SlowOutputStream *)stream)->is_writable = TRUE;

	base_source = g_timeout_source_new (0);

	g_source_set_dummy_callback (base_source);

	pollable_source = g_pollable_source_new (G_OBJECT (stream));

	g_source_add_child_source (pollable_source, base_source);

	g_source_unref (base_source);

	return pollable_source;

}

static void

slow_pollable_output_stream_init (GPollableOutputStreamInterface *pollable_interface,

				  gpointer interface_data)

{

	pollable_interface->is_writable = slow_output_stream_is_writable;

	pollable_interface->write_nonblocking = slow_output_stream_write_nonblocking;

	pollable_interface->create_source = slow_output_stream_create_source;

}

typedef struct {

	GFilterOutputStream parent;

	gboolean is_broken;

} BreakingOutputStream;

typedef struct {

	GFilterOutputStreamClass parent;

} BreakingOutputStreamClass;

GType breaking_output_stream_get_type (void);

static void breaking_pollable_output_stream_init (GPollableOutputStreamInterface *pollable_interface,

						  gpointer interface_data);

G_DEFINE_TYPE_WITH_CODE (BreakingOutputStream, breaking_output_stream,

			 g_type_from_name ("GFilterOutputStream"),

			 G_IMPLEMENT_INTERFACE (G_TYPE_POLLABLE_OUTPUT_STREAM, breaking_pollable_output_stream_init);

			 )

static void

breaking_output_stream_init (BreakingOutputStream *sis)

{

}

static gssize

breaking_output_stream_write (GOutputStream  *stream,

			      const void     *buffer,

			      gsize           count,

			      GCancellable   *cancellable,

			      GError        **error)

{

	if (((BreakingOutputStream *)stream)->is_broken) {

		g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_FAILED, "failed");

		return -1;

	}

	if (count > 128) {

		((BreakingOutputStream *)stream)->is_broken = TRUE;

		count /= 2;

	}

	return g_output_stream_write (G_FILTER_OUTPUT_STREAM (stream)->base_stream,

				      buffer, count, cancellable, error);

}

static void

breaking_output_stream_class_init (BreakingOutputStreamClass *sisclass)

{

	GOutputStreamClass *output_stream_class = G_OUTPUT_STREAM_CLASS (sisclass);

	output_stream_class->write_fn = breaking_output_stream_write;

}

static gboolean

breaking_output_stream_is_writable (GPollableOutputStream *stream)

{

	return TRUE;

}

static gssize

breaking_output_stream_write_nonblocking (GPollableOutputStream  *stream,

					  const void             *buffer,

					  gsize                   count,

					  GError                **error)

{

	if (((BreakingOutputStream *)stream)->is_broken) {

		g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_FAILED, "failed");

		return -1;

	}

	if (count > 128) {

		((BreakingOutputStream *)stream)->is_broken = TRUE;

		count /= 2;

	}

	return g_pollable_output_stream_write_nonblocking (G_POLLABLE_OUTPUT_STREAM (G_FILTER_OUTPUT_STREAM (stream)->base_stream),

							   buffer, count, NULL, error);

}

static GSource *

breaking_output_stream_create_source (GPollableOutputStream *stream,

				      GCancellable *cancellable)

{

	GSource *base_source, *pollable_source;

	base_source = g_timeout_source_new (0);

	g_source_set_dummy_callback (base_source);

	pollable_source = g_pollable_source_new (G_OBJECT (stream));

	g_source_add_child_source (pollable_source, base_source);

	g_source_unref (base_source);

	return pollable_source;

}

static void

breaking_pollable_output_stream_init (GPollableOutputStreamInterface *pollable_interface,

				  gpointer interface_data)

{

	pollable_interface->is_writable = breaking_output_stream_is_writable;

	pollable_interface->write_nonblocking = breaking_output_stream_write_nonblocking;

	pollable_interface->create_source = breaking_output_stream_create_source;

}

#define CHUNK_SIZE 1024

static GString *

chunkify (const char *str, gsize length)

{

	GString *gstr;

	int i, size;

	gstr = g_string_new (NULL);

	for (i = 0; i < length; i += CHUNK_SIZE) {

		size = MIN (CHUNK_SIZE, length - i);

		g_string_append_printf (gstr, "%x\r\n", size);

		g_string_append_len (gstr, str + i, size);

		g_string_append (gstr, "\r\n");

	}

	g_string_append (gstr, "0\r\n\r\n");

	return gstr;

}

static void

do_io_tests (void)

{

	GInputStream *imem, *islow, *in;

	GOutputStream *omem, *oslow, *out;

	GMemoryOutputStream *mem;

	SoupBuffer *raw_contents;

	char *buf;

	GString *chunkified;

	GError *error = NULL;

	gssize nread, nwrote, total;

	gssize chunk_length, chunk_total;

	raw_contents = soup_test_get_index ();

	chunkified = chunkify (raw_contents->data, raw_contents->length);

	debug_printf (1, "  sync read\n");

	imem = g_memory_input_stream_new_from_data (chunkified->str, chunkified->len, NULL);

	islow = g_object_new (slow_input_stream_get_type (),

			      "base-stream", imem,

			      "close-base-stream", TRUE,

			      NULL);

	in = g_object_new (g_type_from_name ("SoupBodyInputStream"),

			   "base-stream", islow,

			   "close-base-stream", TRUE,

			   "encoding", SOUP_ENCODING_CHUNKED,

			   NULL);

	g_object_unref (imem);

	g_object_unref (islow);

	buf = g_malloc (raw_contents->length);

	total = 0;

	while (TRUE) {

		nread = g_input_stream_read (in, buf + total,

					     raw_contents->length - total,

					     NULL, &error);

		g_assert_no_error (error);

		g_clear_error (&error);

		if (nread > 0)

			total += nread;

		else

			break;

	}

	g_input_stream_close (in, NULL, &error);

	g_assert_no_error (error);

	g_clear_error (&error);

	g_object_unref (in);

	soup_assert_cmpmem (buf, total, raw_contents->data, raw_contents->length);

	g_free (buf);

	debug_printf (1, "  async read\n");

	imem = g_memory_input_stream_new_from_data (chunkified->str, chunkified->len, NULL);

	islow = g_object_new (slow_input_stream_get_type (),

			      "base-stream", imem,

			      "close-base-stream", TRUE,

			      NULL);

	in = g_object_new (g_type_from_name ("SoupBodyInputStream"),

			   "base-stream", islow,

			   "close-base-stream", TRUE,

			   "encoding", SOUP_ENCODING_CHUNKED,

			   NULL);

	g_object_unref (imem);

	g_object_unref (islow);

	buf = g_malloc (raw_contents->length);

	total = 0;

	while (TRUE) {

		nread = g_pollable_input_stream_read_nonblocking (G_POLLABLE_INPUT_STREAM (in),

								  buf + total,

								  raw_contents->length - total,

								  NULL, &error);

		if (nread == -1 && g_error_matches (error, G_IO_ERROR, G_IO_ERROR_WOULD_BLOCK)) {

			GSource *source;

			g_clear_error (&error);

			source = g_pollable_input_stream_create_source (G_POLLABLE_INPUT_STREAM (in), NULL);

			g_source_set_dummy_callback (source);

			g_source_attach (source, NULL);

			while (!g_pollable_input_stream_is_readable (G_POLLABLE_INPUT_STREAM (in)))

				g_main_context_iteration (NULL, TRUE);

			g_source_destroy (source);

			g_source_unref (source);

			continue;

		} else if (nread == -1) {

			g_assert_no_error (error);

			g_clear_error (&error);

			break;

		} else if (nread == 0)

			break;

		else

			total += nread;

	}

	g_input_stream_close (in, NULL, &error);

	g_assert_no_error (error);

	g_clear_error (&error);

	g_object_unref (in);

	soup_assert_cmpmem (buf, total, raw_contents->data, raw_contents->length);

	g_free (buf);

	debug_printf (1, "  sync write\n");

	buf = g_malloc (chunkified->len);

	omem = g_memory_output_stream_new (buf, chunkified->len, NULL, NULL);

	oslow = g_object_new (slow_output_stream_get_type (),

			      "base-stream", omem,

			      "close-base-stream", TRUE,

			      NULL);

	out = g_object_new (g_type_from_name ("SoupBodyOutputStream"),

			    "base-stream", oslow,

			    "close-base-stream", TRUE,

			    "encoding", SOUP_ENCODING_CHUNKED,

			    NULL);

	g_object_unref (omem);

	g_object_unref (oslow);

	total = chunk_length = chunk_total = 0;

	while (total < raw_contents->length) {

		if (chunk_total == chunk_length) {

			chunk_length = MIN (CHUNK_SIZE, raw_contents->length - total);

			chunk_total = 0;

		}

		nwrote = g_output_stream_write (out, raw_contents->data + total,

						chunk_length - chunk_total, NULL, &error);

		g_assert_no_error (error);

		g_clear_error (&error);

		if (nwrote > 0) {

			total += nwrote;

			chunk_total += nwrote;

		} else

			break;

	}

	g_output_stream_close (out, NULL, &error);

	g_assert_no_error (error);

	g_clear_error (&error);

	mem = G_MEMORY_OUTPUT_STREAM (omem);

	soup_assert_cmpmem (g_memory_output_stream_get_data (mem),

			    g_memory_output_stream_get_data_size (mem),

			    chunkified->str, chunkified->len);

	g_object_unref (out);

	g_free (buf);

	debug_printf (1, "  async write\n");

	buf = g_malloc (chunkified->len);

	omem = g_memory_output_stream_new (buf, chunkified->len, NULL, NULL);

	oslow = g_object_new (slow_output_stream_get_type (),

			      "base-stream", omem,

			      "close-base-stream", TRUE,

			      NULL);

	out = g_object_new (g_type_from_name ("SoupBodyOutputStream"),

			    "base-stream", oslow,

			    "close-base-stream", TRUE,

			    "encoding", SOUP_ENCODING_CHUNKED,

			    NULL);

	g_object_unref (omem);

	g_object_unref (oslow);

	total = chunk_length = chunk_total = 0;

	while (total < raw_contents->length) {

		if (chunk_total == chunk_length) {

			chunk_length = MIN (CHUNK_SIZE, raw_contents->length - total);

			chunk_total = 0;

		}

		nwrote = g_pollable_output_stream_write_nonblocking (G_POLLABLE_OUTPUT_STREAM (out),

								     raw_contents->data + total,

								     chunk_length - chunk_total,

								     NULL, &error);

		if (nwrote == -1 && g_error_matches (error, G_IO_ERROR, G_IO_ERROR_WOULD_BLOCK)) {

			GSource *source;

			g_clear_error (&error);

			source = g_pollable_output_stream_create_source (G_POLLABLE_OUTPUT_STREAM (out), NULL);

			g_source_set_dummy_callback (source);

			g_source_attach (source, NULL);

			while (!g_pollable_output_stream_is_writable (G_POLLABLE_OUTPUT_STREAM (out)))

				g_main_context_iteration (NULL, TRUE);

			g_source_destroy (source);

			g_source_unref (source);

			continue;

		} else if (nwrote == -1) {

			g_assert_no_error (error);

			g_clear_error (&error);

			break;

		} else {

			total += nwrote;

			chunk_total += nwrote;

		}

	}

	g_output_stream_close (out, NULL, &error);

	g_assert_no_error (error);

	g_clear_error (&error);

	mem = G_MEMORY_OUTPUT_STREAM (omem);

	soup_assert_cmpmem (g_memory_output_stream_get_data (mem),

			    g_memory_output_stream_get_data_size (mem),

			    chunkified->str, chunkified->len);

	g_object_unref (out);

	g_free (buf);

	debug_printf (1, "  failed write\n");

	/* this succeeds if it doesn't critical */

	buf = g_malloc (chunkified->len);

	omem = g_memory_output_stream_new (buf, chunkified->len, NULL, NULL);

	oslow = g_object_new (breaking_output_stream_get_type (),

			      "base-stream", omem,

			      "close-base-stream", TRUE,

			      NULL);

	out = g_object_new (g_type_from_name ("SoupBodyOutputStream"),

			    "base-stream", oslow,

			    "close-base-stream", TRUE,

			    "encoding", SOUP_ENCODING_CHUNKED,

			    NULL);

	g_object_unref (omem);

	g_object_unref (oslow);

	total = 0;

	while (total < raw_contents->length) {

		nwrote = g_output_stream_write (out, raw_contents->data + total,

						raw_contents->length - total, NULL, NULL);

		if (nwrote == -1)

			break;

		else

			total += nwrote;

	}

	g_assert_cmpint (total, !=, raw_contents->length);

	g_output_stream_close (out, NULL, NULL);

	g_object_unref (out);

	g_free (buf);

	g_string_free (chunkified, TRUE);

}

int

main (int argc, char **argv)

{

	int ret;

	test_init (argc, argv, NULL);

	force_io_streams_init ();

	g_test_add_func ("/chunk-io", do_io_tests);

	ret = g_test_run ();

	test_cleanup ();

	return ret;

}