// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Encrypt/decrypt data by xor with a pseudo-random data stream.

package block

import (
	"crypto/block";
	"io";
	"os";
)

// A dataStream is an interface to an unending stream of data,
// used by XorReader and XorWriter to model a pseudo-random generator.
// Calls to Next() return sequential blocks of data from the stream.
// Each call must return at least one byte: there is no EOF.
type dataStream interface {
	Next() []byte
}

type xorReader struct {
	r io.Reader;
	rand dataStream;	// pseudo-random
	buf []byte;		// data available from last call to rand
}

func newXorReader(rand dataStream, r io.Reader) io.Reader {
	x := new(xorReader);
	x.r = r;
	x.rand = rand;
	return x;
}

func (x *xorReader) Read(p []byte) (n int, err os.Error) {
	n, err = x.r.Read(p);

	// xor input with stream.
	bp := 0;
	buf := x.buf;
	for i := 0; i < n; i++ {
		if bp >= len(buf) {
			buf = x.rand.Next();
			bp = 0;
		}
		p[i] ^= buf[bp];
		bp++;
	}
	x.buf = buf[bp:len(buf)];
	return n, err;
}

type xorWriter struct {
	w io.Writer;
	rand dataStream;	// pseudo-random
	buf []byte;	// last buffer returned by rand
	extra []byte;	// extra random data (use before buf)
	work []byte;	// work space
}

func newXorWriter(rand dataStream, w io.Writer) io.Writer {
	x := new(xorWriter);
	x.w = w;
	x.rand = rand;
	x.work = make([]byte, 4096);
	return x;
}

func (x *xorWriter) Write(p []byte) (n int, err os.Error) {
	for len(p) > 0 {
		// Determine next chunk of random data
		// and xor with p into x.work.
		var chunk []byte;
		m := len(p);
		if nn := len(x.extra); nn > 0 {
			// extra points into work, so edit directly
			if m > nn {
				m = nn;
			}
			for i := 0; i < m; i++ {
				x.extra[i] ^= p[i];
			}
			chunk = x.extra[0:m];
		} else {
			// xor p ^ buf into work, refreshing buf as needed
			if nn := len(x.work); m > nn {
				m = nn;
			}
			bp := 0;
			buf := x.buf;
			for i := 0; i < m; i++ {
				if bp >= len(buf) {
					buf = x.rand.Next();
					bp = 0;
				}
				x.work[i] = buf[bp] ^ p[i];
				bp++;
			}
			x.buf = buf[bp:len(buf)];
			chunk = x.work[0:m];
		}

		// Write chunk.
		var nn int;
		nn, err = x.w.Write(chunk);
		if nn != len(chunk) && err == nil {
			err = io.ErrShortWrite;
		}
		if nn < len(chunk) {
			// Reconstruct the random bits from the unwritten
			// data and save them for next time.
			for i := nn; i < m; i++ {
				chunk[i] ^= p[i];
			}
			x.extra = chunk[nn:len(chunk)];
		}
		n += nn;
		if err != nil {
			return;
		}
		p = p[m:len(p)];
	}
	return;
}