~samwhited/xmpp

xmpp/compress/methods.go -rw-r--r-- 3.0 KiB
c9743d9fSam Whited docs: update changelog to include deprecations 3 days ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
// Copyright 2016 The Mellium Contributors.
// Use of this source code is governed by the BSD 2-clause
// license that can be found in the LICENSE file.

package compress

import (
	"compress/lzw"
	"compress/zlib"
	"io"
	"sync"
)

var (
	// LZW implements stream compression using the Lempel-Ziv-Welch (DCLZ)
	// compressed data format.
	LZW Method = lzwMethod
)

var lzwMethod = Method{
	Name: "lzw",
	Wrapper: func(rw io.ReadWriter) (io.ReadWriter, error) {
		rc := lzw.NewReader(rw, lzw.LSB, 8)
		wc := lzw.NewWriter(rw, lzw.LSB, 8)
		return struct {
			io.Reader
			io.Writer
			io.Closer
		}{
			Reader: rc,
			Writer: wc,
			Closer: &multiCloser{rc, wc},
		}, nil
	},
}

// Method is a stream compression method.
// Custom methods may be defined, but generally speaking the only supported
// methods will be those with names defined in the "Stream Compression Methods
// Registry" maintained by the XSF Editor:
// https://xmpp.org/registrar/compress.html
//
// Since ZLIB is always supported, a Method is not defined for it.
type Method struct {
	Name    string
	Wrapper func(io.ReadWriter) (io.ReadWriter, error)
}

type multiCloser []io.Closer

// Close attempts to call every close method in the multiCloser.
// It always attempts all of them (unless one of them panics), but it only
// returns the last error if multiple of them error.
// There's probably a better way.
func (mc multiCloser) Close() (err error) {
	var e error
	for _, c := range mc {
		// Return one error; we don't really care what.
		if e = c.Close(); e != nil {
			err = e
		}
	}
	return err
}

// zlibDelayedSetup is an io.ReadWriteCloser that uses an underlying zlib reader
// and writer, but defers creation of the reader until the first read. It is a
// jank hack to work around the fact that the zlib reader tries to immediately
// read header data from the connection (and blocks until it can do so), but if
// we're a client we need to go ahead and return the new reader and send a
// <stream:stream> before reading the header data.
type zlibDelayedSetup struct {
	wm, rm sync.Mutex

	raw        io.ReadWriter
	zlibWriter *zlib.Writer
	zlibReader io.ReadCloser
}

func (r *zlibDelayedSetup) readSetup() (err error) {
	r.rm.Lock()
	defer r.rm.Unlock()

	if r.zlibReader == nil {
		if r.zlibReader, err = zlib.NewReader(r.raw); err != nil {
			return
		}
	}

	return
}

func (r *zlibDelayedSetup) Write(p []byte) (n int, err error) {
	if n, err = r.zlibWriter.Write(p); err != nil {
		return
	}
	return n, r.zlibWriter.Flush()
}

func (r *zlibDelayedSetup) Read(p []byte) (n int, err error) {
	if err = r.readSetup(); err != nil {
		return
	}
	return r.zlibReader.Read(p)
}

func (r *zlibDelayedSetup) Close() error {

	mc := multiCloser{}

	r.rm.Lock()
	defer r.rm.Unlock()
	if r.zlibReader != nil {
		mc = append(mc, r.zlibReader)
	}

	r.wm.Lock()
	defer r.wm.Unlock()
	if r.zlibWriter != nil {
		mc = append(mc, r.zlibWriter)
	}

	return mc.Close()
}

var zlibMethod = Method{
	Name: "zlib",
	Wrapper: func(rw io.ReadWriter) (io.ReadWriter, error) {
		return &zlibDelayedSetup{raw: rw, zlibWriter: zlib.NewWriter(rw)}, nil
	},
}