1279801Smarkj/*-
2279801Smarkj * Copyright (c) 2014 Mark Johnston <markj@FreeBSD.org>
3204552Salfred *
4279801Smarkj * Redistribution and use in source and binary forms, with or without
5279801Smarkj * modification, are permitted provided that the following conditions are
6279801Smarkj * met:
7279801Smarkj * 1. Redistributions of source code must retain the above copyright
8279801Smarkj *    notice, this list of conditions and the following disclaimer.
9279801Smarkj * 2. Redistributions in binary form must reproduce the above copyright
10279801Smarkj *    notice, this list of conditions and the following disclaimer in
11279801Smarkj *    the documentation and/or other materials provided with the
12279801Smarkj *    distribution.
13204552Salfred *
14279801Smarkj * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
15279801Smarkj * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
16279801Smarkj * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
17279801Smarkj * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
18279801Smarkj * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
19279801Smarkj * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
20279801Smarkj * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
21279801Smarkj * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
22279801Smarkj * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
23279801Smarkj * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
24279801Smarkj * SUCH DAMAGE.
25204552Salfred */
26204552Salfred
27279801Smarkj#include <sys/cdefs.h>
28279801Smarkj__FBSDID("$FreeBSD$");
29204552Salfred
30279801Smarkj#include <sys/param.h>
31204552Salfred
32279801Smarkj#include <sys/gzio.h>
33279801Smarkj#include <sys/kernel.h>
34204552Salfred#include <sys/malloc.h>
35281855Srodrigc#include <sys/zutil.h>
36279801Smarkj
37279801Smarkj#define	KERN_GZ_HDRLEN		10	/* gzip header length */
38279801Smarkj#define	KERN_GZ_TRAILERLEN	8	/* gzip trailer length */
39279801Smarkj#define	KERN_GZ_MAGIC1		0x1f	/* first magic byte */
40279801Smarkj#define	KERN_GZ_MAGIC2		0x8b	/* second magic byte */
41204552Salfred
42279801SmarkjMALLOC_DEFINE(M_GZIO, "gzio", "zlib state");
43204552Salfred
44279801Smarkjstruct gzio_stream {
45279801Smarkj	uint8_t *	gz_buffer;	/* output buffer */
46279801Smarkj	size_t		gz_bufsz;	/* total buffer size */
47279801Smarkj	off_t		gz_off;		/* offset into the output stream */
48279801Smarkj	enum gzio_mode	gz_mode;	/* stream mode */
49279801Smarkj	uint32_t	gz_crc;		/* stream CRC32 */
50279801Smarkj	gzio_cb		gz_cb;		/* output callback */
51279801Smarkj	void *		gz_arg;		/* private callback arg */
52279801Smarkj	z_stream	gz_stream;	/* zlib state */
53279801Smarkj};
54204552Salfred
55279801Smarkjstatic void *	gz_alloc(void *, u_int, u_int);
56279801Smarkjstatic void	gz_free(void *, void *);
57279801Smarkjstatic int	gz_write(struct gzio_stream *, void *, u_int, int);
58204552Salfred
59279801Smarkjstruct gzio_stream *
60279801Smarkjgzio_init(gzio_cb cb, enum gzio_mode mode, size_t bufsz, int level, void *arg)
61204552Salfred{
62279801Smarkj	struct gzio_stream *s;
63279801Smarkj	uint8_t *hdr;
64279801Smarkj	int error;
65204552Salfred
66279801Smarkj	if (bufsz < KERN_GZ_HDRLEN)
67279801Smarkj		return (NULL);
68279801Smarkj	if (mode != GZIO_DEFLATE)
69279801Smarkj		return (NULL);
70204552Salfred
71279801Smarkj	s = gz_alloc(NULL, 1, sizeof(*s));
72279801Smarkj	s->gz_bufsz = bufsz;
73279801Smarkj	s->gz_buffer = gz_alloc(NULL, 1, s->gz_bufsz);
74279801Smarkj	s->gz_mode = mode;
75279801Smarkj	s->gz_crc = ~0U;
76279801Smarkj	s->gz_cb = cb;
77279801Smarkj	s->gz_arg = arg;
78204552Salfred
79279801Smarkj	s->gz_stream.zalloc = gz_alloc;
80279801Smarkj	s->gz_stream.zfree = gz_free;
81279801Smarkj	s->gz_stream.opaque = NULL;
82279801Smarkj	s->gz_stream.next_in = Z_NULL;
83279801Smarkj	s->gz_stream.avail_in = 0;
84204552Salfred
85279801Smarkj	error = deflateInit2(&s->gz_stream, level, Z_DEFLATED, -MAX_WBITS,
86279801Smarkj	    DEF_MEM_LEVEL, Z_DEFAULT_STRATEGY);
87279801Smarkj	if (error != 0)
88279801Smarkj		goto fail;
89204552Salfred
90279801Smarkj	s->gz_stream.avail_out = s->gz_bufsz;
91279801Smarkj	s->gz_stream.next_out = s->gz_buffer;
92204552Salfred
93279801Smarkj	/* Write the gzip header to the output buffer. */
94279801Smarkj	hdr = s->gz_buffer;
95279801Smarkj	memset(hdr, 0, KERN_GZ_HDRLEN);
96279801Smarkj	hdr[0] = KERN_GZ_MAGIC1;
97279801Smarkj	hdr[1] = KERN_GZ_MAGIC2;
98279801Smarkj	hdr[2] = Z_DEFLATED;
99279801Smarkj	hdr[9] = OS_CODE;
100279801Smarkj	s->gz_stream.next_out += KERN_GZ_HDRLEN;
101279801Smarkj	s->gz_stream.avail_out -= KERN_GZ_HDRLEN;
102204552Salfred
103279801Smarkj	return (s);
104204552Salfred
105279801Smarkjfail:
106279801Smarkj	gz_free(NULL, s->gz_buffer);
107279801Smarkj	gz_free(NULL, s);
108279801Smarkj	return (NULL);
109279801Smarkj}
110204552Salfred
111279801Smarkjint
112279801Smarkjgzio_write(struct gzio_stream *s, void *data, u_int len)
113279801Smarkj{
114204552Salfred
115279801Smarkj	return (gz_write(s, data, len, Z_NO_FLUSH));
116204552Salfred}
117204552Salfred
118279801Smarkjint
119279801Smarkjgzio_flush(struct gzio_stream *s)
120204552Salfred{
121204552Salfred
122279801Smarkj	return (gz_write(s, NULL, 0, Z_FINISH));
123204552Salfred}
124204552Salfred
125279801Smarkjvoid
126279801Smarkjgzio_fini(struct gzio_stream *s)
127204552Salfred{
128204552Salfred
129279801Smarkj	(void)deflateEnd(&s->gz_stream);
130279801Smarkj	gz_free(NULL, s->gz_buffer);
131279801Smarkj	gz_free(NULL, s);
132204552Salfred}
133204552Salfred
134279801Smarkjstatic void *
135279801Smarkjgz_alloc(void *arg __unused, u_int n, u_int sz)
136204552Salfred{
137204552Salfred
138279801Smarkj	/*
139279801Smarkj	 * Memory for zlib state is allocated using M_NODUMP since it may be
140279801Smarkj	 * used to compress a kernel dump, and we don't want zlib to attempt to
141279801Smarkj	 * compress its own state.
142279801Smarkj	 */
143279801Smarkj	return (malloc(n * sz, M_GZIO, M_WAITOK | M_ZERO | M_NODUMP));
144204552Salfred}
145204552Salfred
146279801Smarkjstatic void
147279801Smarkjgz_free(void *arg __unused, void *ptr)
148204552Salfred{
149204552Salfred
150279801Smarkj	free(ptr, M_GZIO);
151204552Salfred}
152204552Salfred
153279801Smarkjstatic int
154279801Smarkjgz_write(struct gzio_stream *s, void *buf, u_int len, int zflag)
155204552Salfred{
156279801Smarkj	uint8_t trailer[KERN_GZ_TRAILERLEN];
157279801Smarkj	size_t room;
158279801Smarkj	int error, zerror;
159204552Salfred
160279801Smarkj	KASSERT(zflag == Z_FINISH || zflag == Z_NO_FLUSH,
161279801Smarkj	    ("unexpected flag %d", zflag));
162279801Smarkj	KASSERT(s->gz_mode == GZIO_DEFLATE,
163279801Smarkj	    ("invalid stream mode %d", s->gz_mode));
164204552Salfred
165279801Smarkj	if (len > 0) {
166279801Smarkj		s->gz_stream.avail_in = len;
167279801Smarkj		s->gz_stream.next_in = buf;
168279801Smarkj		s->gz_crc = crc32_raw(buf, len, s->gz_crc);
169279801Smarkj	} else
170279801Smarkj		s->gz_crc ^= ~0U;
171204552Salfred
172279801Smarkj	error = 0;
173279801Smarkj	do {
174279801Smarkj		zerror = deflate(&s->gz_stream, zflag);
175279801Smarkj		if (zerror != Z_OK && zerror != Z_STREAM_END) {
176279801Smarkj			error = EIO;
177279801Smarkj			break;
178279801Smarkj		}
179204552Salfred
180279801Smarkj		if (s->gz_stream.avail_out == 0 || zerror == Z_STREAM_END) {
181279801Smarkj			/*
182279801Smarkj			 * Our output buffer is full or there's nothing left
183279801Smarkj			 * to produce, so we're flushing the buffer.
184279801Smarkj			 */
185279801Smarkj			len = s->gz_bufsz - s->gz_stream.avail_out;
186279801Smarkj			if (zerror == Z_STREAM_END) {
187279801Smarkj				/*
188279801Smarkj				 * Try to pack as much of the trailer into the
189279801Smarkj				 * output buffer as we can.
190279801Smarkj				 */
191279801Smarkj				((uint32_t *)trailer)[0] = s->gz_crc;
192279801Smarkj				((uint32_t *)trailer)[1] =
193279801Smarkj				    s->gz_stream.total_in;
194279801Smarkj				room = MIN(KERN_GZ_TRAILERLEN,
195279801Smarkj				    s->gz_bufsz - len);
196279801Smarkj				memcpy(s->gz_buffer + len, trailer, room);
197279801Smarkj				len += room;
198279801Smarkj			}
199204552Salfred
200279801Smarkj			error = s->gz_cb(s->gz_buffer, len, s->gz_off,
201279801Smarkj			    s->gz_arg);
202279801Smarkj			if (error != 0)
203279801Smarkj				break;
204204552Salfred
205279801Smarkj			s->gz_off += len;
206279801Smarkj			s->gz_stream.next_out = s->gz_buffer;
207279801Smarkj			s->gz_stream.avail_out = s->gz_bufsz;
208204552Salfred
209279801Smarkj			/*
210279801Smarkj			 * If we couldn't pack the trailer into the output
211279801Smarkj			 * buffer, write it out now.
212279801Smarkj			 */
213279801Smarkj			if (zerror == Z_STREAM_END && room < KERN_GZ_TRAILERLEN)
214279801Smarkj				error = s->gz_cb(trailer + room,
215279801Smarkj				    KERN_GZ_TRAILERLEN - room, s->gz_off,
216279801Smarkj				    s->gz_arg);
217279801Smarkj		}
218279801Smarkj	} while (zerror != Z_STREAM_END &&
219279801Smarkj	    (zflag == Z_FINISH || s->gz_stream.avail_in > 0));
220279801Smarkj
221279801Smarkj	return (error);
222204552Salfred}
223