go-tg-screenshot-bot

1
package xgb
2

3
import (
4
	"errors"
5
	"io"
6
	"log"
7
	"net"
8
	"os"
9
	"sync"
10
)
11

12
var (
13
	// Where to log error-messages. Defaults to stderr.
14
	// To disable logging, just set this to log.New(ioutil.Discard, "", 0)
15
	Logger = log.New(os.Stderr, "XGB: ", log.Lshortfile)
16
)
17

18
const (
19
	// cookieBuffer represents the queue size of cookies existing at any
20
	// point in time. The size of the buffer is really only important when
21
	// there are many requests without replies made in sequence. Once the
22
	// buffer fills, a round trip request is made to clear the buffer.
23
	cookieBuffer = 1000
24

25
	// xidBuffer represents the queue size of the xid channel.
26
	// I don't think this value matters much, since xid generation is not
27
	// that expensive.
28
	xidBuffer = 5
29

30
	// seqBuffer represents the queue size of the sequence number channel.
31
	// I don't think this value matters much, since sequence number generation
32
	// is not that expensive.
33
	seqBuffer = 5
34

35
	// reqBuffer represents the queue size of the number of requests that
36
	// can be made until new ones block. This value seems OK.
37
	reqBuffer = 100
38

39
	// eventBuffer represents the queue size of the number of events or errors
40
	// that can be loaded off the wire and not grabbed with WaitForEvent
41
	// until reading an event blocks. This value should be big enough to handle
42
	// bursts of events.
43
	eventBuffer = 5000
44
)
45

46
// A Conn represents a connection to an X server.
47
type Conn struct {
48
	host          string
49
	conn          net.Conn
50
	display       string
51
	DisplayNumber int
52
	DefaultScreen int
53
	SetupBytes    []byte
54

55
	setupResourceIdBase uint32
56
	setupResourceIdMask uint32
57

58
	eventChan  chan eventOrError
59
	cookieChan chan *Cookie
60
	xidChan    chan xid
61
	seqChan    chan uint16
62
	reqChan    chan *request
63
	doneSend   chan struct{}
64
	doneRead   chan struct{}
65

66
	// ExtLock is a lock used whenever new extensions are initialized.
67
	// It should not be used. It is exported for use in the extension
68
	// sub-packages.
69
	ExtLock sync.RWMutex
70

71
	// Extensions is a map from extension name to major opcode. It should
72
	// not be used. It is exported for use in the extension sub-packages.
73
	Extensions map[string]byte
74
}
75

76
// NewConn creates a new connection instance. It initializes locks, data
77
// structures, and performs the initial handshake. (The code for the handshake
78
// has been relegated to conn.go.)
79
// It is up to user to close connection with Close() method to finish all unfinished requests and clean up spawned goroutines.
80
// If the connection unexpectedly closes itself and WaitForEvent() returns "nil, nil", everything is cleaned by that moment, but nothing bad happens if you call Close() after.
81
func NewConn() (*Conn, error) {
82
	return NewConnDisplay("")
83
}
84

85
// NewConnDisplay is just like NewConn (see closing instructions), but allows a specific DISPLAY
86
// string to be used.
87
// If 'display' is empty it will be taken from os.Getenv("DISPLAY").
88
//
89
// Examples:
90
//	NewConn(":1") -> net.Dial("unix", "", "/tmp/.X11-unix/X1")
91
//	NewConn("/tmp/launch-12/:0") -> net.Dial("unix", "", "/tmp/launch-12/:0")
92
//	NewConn("hostname:2.1") -> net.Dial("tcp", "", "hostname:6002")
93
//	NewConn("tcp/hostname:1.0") -> net.Dial("tcp", "", "hostname:6001")
94
func NewConnDisplay(display string) (*Conn, error) {
95
	c := &Conn{}
96

97
	// First connect. This reads authority, checks DISPLAY environment
98
	// variable, and loads the initial Setup info.
99
	err := c.connect(display)
100
	if err != nil {
101
		return nil, err
102
	}
103

104
	return postNewConn(c)
105
}
106

107
// NewConnNet is just like NewConn (see closing instructions), but allows a specific net.Conn
108
// to be used.
109
func NewConnNet(netConn net.Conn) (*Conn, error) {
110
	c := &Conn{}
111

112
	// First connect. This reads authority, checks DISPLAY environment
113
	// variable, and loads the initial Setup info.
114
	err := c.connectNet(netConn)
115

116
	if err != nil {
117
		return nil, err
118
	}
119

120
	return postNewConn(c)
121
}
122

123
func postNewConn(c *Conn) (*Conn, error) {
124
	c.Extensions = make(map[string]byte)
125

126
	c.cookieChan = make(chan *Cookie, cookieBuffer)
127
	c.xidChan = make(chan xid, xidBuffer)
128
	c.seqChan = make(chan uint16, seqBuffer)
129
	c.reqChan = make(chan *request, reqBuffer)
130
	c.eventChan = make(chan eventOrError, eventBuffer)
131
	c.doneSend = make(chan struct{})
132
	c.doneRead = make(chan struct{})
133

134
	go c.generateXIds()
135
	go c.generateSeqIds()
136
	go c.sendRequests()
137
	go c.readResponses()
138

139
	return c, nil
140
}
141

142
// Close gracefully closes the connection to the X server.
143
// When everything is cleaned up, the WaitForEvent method will return (nil, nil)
144
func (c *Conn) Close() {
145
	select {
146
	case c.reqChan <- nil:
147
	case <-c.doneSend:
148
	}
149
}
150

151
// Event is an interface that can contain any of the events returned by the
152
// server. Use a type assertion switch to extract the Event structs.
153
type Event interface {
154
	Bytes() []byte
155
	String() string
156
}
157

158
// NewEventFun is the type of function use to construct events from raw bytes.
159
// It should not be used. It is exported for use in the extension sub-packages.
160
type NewEventFun func(buf []byte) Event
161

162
// NewEventFuncs is a map from event numbers to functions that create
163
// the corresponding event. It should not be used. It is exported for use
164
// in the extension sub-packages.
165
var NewEventFuncs = make(map[int]NewEventFun)
166

167
// NewExtEventFuncs is a temporary map that stores event constructor functions
168
// for each extension. When an extension is initialized, each event for that
169
// extension is added to the 'NewEventFuncs' map. It should not be used. It is
170
// exported for use in the extension sub-packages.
171
var NewExtEventFuncs = make(map[string]map[int]NewEventFun)
172

173
// Error is an interface that can contain any of the errors returned by
174
// the server. Use a type assertion switch to extract the Error structs.
175
type Error interface {
176
	SequenceId() uint16
177
	BadId() uint32
178
	Error() string
179
}
180

181
// NewErrorFun is the type of function use to construct errors from raw bytes.
182
// It should not be used. It is exported for use in the extension sub-packages.
183
type NewErrorFun func(buf []byte) Error
184

185
// NewErrorFuncs is a map from error numbers to functions that create
186
// the corresponding error. It should not be used. It is exported for use in
187
// the extension sub-packages.
188
var NewErrorFuncs = make(map[int]NewErrorFun)
189

190
// NewExtErrorFuncs is a temporary map that stores error constructor functions
191
// for each extension. When an extension is initialized, each error for that
192
// extension is added to the 'NewErrorFuncs' map. It should not be used. It is
193
// exported for use in the extension sub-packages.
194
var NewExtErrorFuncs = make(map[string]map[int]NewErrorFun)
195

196
// eventOrError corresponds to values that can be either an event or an
197
// error.
198
type eventOrError interface{}
199

200
// NewId generates a new unused ID for use with requests like CreateWindow.
201
// If no new ids can be generated, the id returned is 0 and error is non-nil.
202
// This shouldn't be used directly, and is exported for use in the extension
203
// sub-packages.
204
// If you need identifiers, use the appropriate constructor.
205
// e.g., For a window id, use xproto.NewWindowId. For
206
// a new pixmap id, use xproto.NewPixmapId. And so on.
207
// Returns (0, io.EOF) when the connection is closed.
208
func (c *Conn) NewId() (uint32, error) {
209
	xid, ok := <-c.xidChan
210
	if !ok {
211
		return 0, io.EOF
212
	}
213
	if xid.err != nil {
214
		return 0, xid.err
215
	}
216
	return xid.id, nil
217
}
218

219
// xid encapsulates a resource identifier being sent over the Conn.xidChan
220
// channel. If no new resource id can be generated, id is set to 0 and a
221
// non-nil error is set in xid.err.
222
type xid struct {
223
	id  uint32
224
	err error
225
}
226

227
// generateXids sends new Ids down the channel for NewId to use.
228
// generateXids should be run in its own goroutine.
229
// This needs to be updated to use the XC Misc extension once we run out of
230
// new ids.
231
// Thanks to libxcb/src/xcb_xid.c. This code is greatly inspired by it.
232
func (c *Conn) generateXIds() {
233
	defer close(c.xidChan)
234

235
	// This requires some explanation. From the horse's mouth:
236
	// "The resource-id-mask contains a single contiguous set of bits (at least
237
	// 18).  The client allocates resource IDs for types WINDOW, PIXMAP,
238
	// CURSOR, FONT, GCONTEXT, and COLORMAP by choosing a value with only some
239
	// subset of these bits set and ORing it with resource-id-base. Only values
240
	// constructed in this way can be used to name newly created resources over
241
	// this connection."
242
	// So for example (using 8 bit integers), the mask might look like:
243
	// 00111000
244
	// So that valid values would be 00101000, 00110000, 00001000, and so on.
245
	// Thus, the idea is to increment it by the place of the last least
246
	// significant '1'. In this case, that value would be 00001000. To get
247
	// that value, we can AND the original mask with its two's complement:
248
	// 00111000 & 11001000 = 00001000.
249
	// And we use that value to increment the last resource id to get a new one.
250
	// (And then, of course, we OR it with resource-id-base.)
251
	inc := c.setupResourceIdMask & -c.setupResourceIdMask
252
	max := c.setupResourceIdMask
253
	last := uint32(0)
254
	for {
255
		id := xid{}
256
		if last > 0 && last >= max-inc+1 {
257
			// TODO: Use the XC Misc extension to look for released ids.
258
			id = xid{
259
				id:  0,
260
				err: errors.New("There are no more available resource identifiers."),
261
			}
262
		} else {
263
			last += inc
264
			id = xid{
265
				id:  last | c.setupResourceIdBase,
266
				err: nil,
267
			}
268
		}
269

270
		select {
271
		case c.xidChan <- id:
272
		case <-c.doneSend:
273
			// c.sendRequests is down and since this id is used by requests, we don't need this goroutine running anymore.
274
			return
275
		}
276
	}
277
}
278

279
// newSeqId fetches the next sequence id from the Conn.seqChan channel.
280
func (c *Conn) newSequenceId() uint16 {
281
	return <-c.seqChan
282
}
283

284
// generateSeqIds returns new sequence ids. It is meant to be run in its
285
// own goroutine.
286
// A sequence id is generated for *every* request. It's the identifier used
287
// to match up replies with requests.
288
// Since sequence ids can only be 16 bit integers we start over at zero when it
289
// comes time to wrap.
290
// N.B. As long as the cookie buffer is less than 2^16, there are no limitations
291
// on the number (or kind) of requests made in sequence.
292
func (c *Conn) generateSeqIds() {
293
	defer close(c.seqChan)
294

295
	seqid := uint16(1)
296
	for {
297
		select {
298
		case c.seqChan <- seqid:
299
			if seqid == uint16((1<<16)-1) {
300
				seqid = 0
301
			} else {
302
				seqid++
303
			}
304
		case <-c.doneSend:
305
			// c.sendRequests is down and since only that function uses sequence ids (via newSequenceId method), we don't need this goroutine running anymore.
306
			return
307
		}
308
	}
309
}
310

311
// request encapsulates a buffer of raw bytes (containing the request data)
312
// and a cookie, which when combined represents a single request.
313
// The cookie is used to match up the reply/error.
314
type request struct {
315
	buf    []byte
316
	cookie *Cookie
317

318
	// seq is closed when the request (cookie) has been sequenced by the Conn.
319
	seq chan struct{}
320
}
321

322
// NewRequest takes the bytes and a cookie of a particular request, constructs
323
// a request type, and sends it over the Conn.reqChan channel.
324
// Note that the sequence number is added to the cookie after it is sent
325
// over the request channel, but before it is sent to X.
326
//
327
// Note that you may safely use NewRequest to send arbitrary byte requests
328
// to X. The resulting cookie can be used just like any normal cookie and
329
// abides by the same rules, except that for replies, you'll get back the
330
// raw byte data. This may be useful for performance critical sections where
331
// every allocation counts, since all X requests in XGB allocate a new byte
332
// slice. In contrast, NewRequest allocates one small request struct and
333
// nothing else. (Except when the cookie buffer is full and has to be flushed.)
334
//
335
// If you're using NewRequest manually, you'll need to use NewCookie to create
336
// a new cookie.
337
//
338
// In all likelihood, you should be able to copy and paste with some minor
339
// edits the generated code for the request you want to issue.
340
func (c *Conn) NewRequest(buf []byte, cookie *Cookie) {
341
	seq := make(chan struct{})
342
	select {
343
	case c.reqChan <- &request{buf: buf, cookie: cookie, seq: seq}:
344
		// request is in buffer
345
		// wait until request is processed or connection is closed
346
		select {
347
		case <-seq:
348
			// request was successfully sent to X server
349
		case <-c.doneSend:
350
			// c.sendRequests is down, your request was not handled
351
		}
352
	case <-c.doneSend:
353
		// c.sendRequests is down, nobody is listening to your requests
354
	}
355
}
356

357
// sendRequests is run as a single goroutine that takes requests and writes
358
// the bytes to the wire and adds the cookie to the cookie queue.
359
// It is meant to be run as its own goroutine.
360
func (c *Conn) sendRequests() {
361
	defer close(c.cookieChan)
362
	defer c.conn.Close()
363
	defer close(c.doneSend)
364

365
	for {
366
		select {
367
		case req := <-c.reqChan:
368
			if req == nil {
369
				// a request by c.Close() to gracefully exit
370
				// Flush the response reading goroutine.
371
				if err := c.noop(); err != nil {
372
					c.conn.Close()
373
					<-c.doneRead
374
				}
375
				return
376
			}
377
			// ho there! if the cookie channel is nearly full, force a round
378
			// trip to clear out the cookie buffer.
379
			// Note that we circumvent the request channel, because we're *in*
380
			// the request channel.
381
			if len(c.cookieChan) == cookieBuffer-1 {
382
				if err := c.noop(); err != nil {
383
					// Shut everything down.
384
					c.conn.Close()
385
					<-c.doneRead
386
					return
387
				}
388
			}
389
			req.cookie.Sequence = c.newSequenceId()
390
			c.cookieChan <- req.cookie
391
			if err := c.writeBuffer(req.buf); err != nil {
392
				c.conn.Close()
393
				<-c.doneRead
394
				return
395
			}
396
			close(req.seq)
397
		case <-c.doneRead:
398
			return
399
		}
400
	}
401
}
402

403
// noop circumvents the usual request sending goroutines and forces a round
404
// trip request manually.
405
func (c *Conn) noop() error {
406
	cookie := c.NewCookie(true, true)
407
	cookie.Sequence = c.newSequenceId()
408
	c.cookieChan <- cookie
409
	if err := c.writeBuffer(c.getInputFocusRequest()); err != nil {
410
		return err
411
	}
412
	cookie.Reply() // wait for the buffer to clear
413
	return nil
414
}
415

416
// writeBuffer is a convenience function for writing a byte slice to the wire.
417
func (c *Conn) writeBuffer(buf []byte) error {
418
	if _, err := c.conn.Write(buf); err != nil {
419
		Logger.Printf("A write error is unrecoverable: %s", err)
420
		return err
421
	}
422
	return nil
423
}
424

425
// readResponses is a goroutine that reads events, errors and
426
// replies off the wire.
427
// When an event is read, it is always added to the event channel.
428
// When an error is read, if it corresponds to an existing checked cookie,
429
// it is sent to that cookie's error channel. Otherwise it is added to the
430
// event channel.
431
// When a reply is read, it is added to the corresponding cookie's reply
432
// channel. (It is an error if no such cookie exists in this case.)
433
// Finally, cookies that came "before" this reply are always cleaned up.
434
func (c *Conn) readResponses() {
435
	defer close(c.eventChan)
436
	defer c.conn.Close()
437
	defer close(c.doneRead)
438

439
	var (
440
		err        Error
441
		seq        uint16
442
		replyBytes []byte
443
	)
444

445
	for {
446
		buf := make([]byte, 32)
447
		err, seq = nil, 0
448
		if _, err := io.ReadFull(c.conn, buf); err != nil {
449
			select {
450
			case <-c.doneSend:
451
				// gracefully closing
452
				return
453
			default:
454
			}
455
			Logger.Printf("A read error is unrecoverable: %s", err)
456
			c.eventChan <- err
457
			return
458
		}
459
		switch buf[0] {
460
		case 0: // This is an error
461
			// Use the constructor function for this error (that is auto
462
			// generated) by looking it up by the error number.
463
			newErrFun, ok := NewErrorFuncs[int(buf[1])]
464
			if !ok {
465
				Logger.Printf("BUG: Could not find error constructor function "+
466
					"for error with number %d.", buf[1])
467
				continue
468
			}
469
			err = newErrFun(buf)
470
			seq = err.SequenceId()
471

472
			// This error is either sent to the event channel or a specific
473
			// cookie's error channel below.
474
		case 1: // This is a reply
475
			seq = Get16(buf[2:])
476

477
			// check to see if this reply has more bytes to be read
478
			size := Get32(buf[4:])
479
			if size > 0 {
480
				byteCount := 32 + size*4
481
				biggerBuf := make([]byte, byteCount)
482
				copy(biggerBuf[:32], buf)
483
				if _, err := io.ReadFull(c.conn, biggerBuf[32:]); err != nil {
484
					Logger.Printf("A read error is unrecoverable: %s", err)
485
					c.eventChan <- err
486
					return
487
				}
488
				replyBytes = biggerBuf
489
			} else {
490
				replyBytes = buf
491
			}
492

493
			// This reply is sent to its corresponding cookie below.
494
		default: // This is an event
495
			// Use the constructor function for this event (like for errors,
496
			// and is also auto generated) by looking it up by the event number.
497
			// Note that we AND the event number with 127 so that we ignore
498
			// the most significant bit (which is set when it was sent from
499
			// a SendEvent request).
500
			evNum := int(buf[0] & 127)
501
			newEventFun, ok := NewEventFuncs[evNum]
502
			if !ok {
503
				Logger.Printf("BUG: Could not find event construct function "+
504
					"for event with number %d.", evNum)
505
				continue
506
			}
507
			c.eventChan <- newEventFun(buf)
508
			continue
509
		}
510

511
		// At this point, we have a sequence number and we're either
512
		// processing an error or a reply, which are both responses to
513
		// requests. So all we have to do is find the cookie corresponding
514
		// to this error/reply, and send the appropriate data to it.
515
		// In doing so, we make sure that any cookies that came before it
516
		// are marked as successful if they are void and checked.
517
		// If there's a cookie that requires a reply that is before this
518
		// reply, then something is wrong.
519
		for cookie := range c.cookieChan {
520
			// This is the cookie we're looking for. Process and break.
521
			if cookie.Sequence == seq {
522
				if err != nil { // this is an error to a request
523
					// synchronous processing
524
					if cookie.errorChan != nil {
525
						cookie.errorChan <- err
526
					} else { // asynchronous processing
527
						c.eventChan <- err
528
						// if this is an unchecked reply, ping the cookie too
529
						if cookie.pingChan != nil {
530
							cookie.pingChan <- true
531
						}
532
					}
533
				} else { // this is a reply
534
					if cookie.replyChan == nil {
535
						Logger.Printf("Reply with sequence id %d does not "+
536
							"have a cookie with a valid reply channel.", seq)
537
						continue
538
					} else {
539
						cookie.replyChan <- replyBytes
540
					}
541
				}
542
				break
543
			}
544

545
			switch {
546
			// Checked requests with replies
547
			case cookie.replyChan != nil && cookie.errorChan != nil:
548
				Logger.Printf("Found cookie with sequence id %d that is "+
549
					"expecting a reply but will never get it. Currently "+
550
					"on sequence number %d", cookie.Sequence, seq)
551
			// Unchecked requests with replies
552
			case cookie.replyChan != nil && cookie.pingChan != nil:
553
				Logger.Printf("Found cookie with sequence id %d that is "+
554
					"expecting a reply (and not an error) but will never "+
555
					"get it. Currently on sequence number %d",
556
					cookie.Sequence, seq)
557
			// Checked requests without replies
558
			case cookie.pingChan != nil && cookie.errorChan != nil:
559
				cookie.pingChan <- true
560
				// Unchecked requests without replies don't have any channels,
561
				// so we can't do anything with them except let them pass by.
562
			}
563
		}
564
	}
565
}
566

567
// processEventOrError takes an eventOrError, type switches on it,
568
// and returns it in Go idiomatic style.
569
func processEventOrError(everr eventOrError) (Event, Error) {
570
	switch ee := everr.(type) {
571
	case Event:
572
		return ee, nil
573
	case Error:
574
		return nil, ee
575
	case error:
576
		// c.conn read error
577
	case nil:
578
		// c.eventChan is closed
579
	default:
580
		Logger.Printf("Invalid event/error type: %T", everr)
581
	}
582
	return nil, nil
583
}
584

585
// WaitForEvent returns the next event from the server.
586
// It will block until an event is available.
587
// WaitForEvent returns either an Event or an Error. (Returning both
588
// is a bug.) Note than an Error here is an X error and not an XGB error. That
589
// is, X errors are sometimes completely expected (and you may want to ignore
590
// them in some cases).
591
//
592
// If both the event and error are nil, then the connection has been closed.
593
func (c *Conn) WaitForEvent() (Event, Error) {
594
	return processEventOrError(<-c.eventChan)
595
}
596

597
// PollForEvent returns the next event from the server if one is available in
598
// the internal queue without blocking. Note that unlike WaitForEvent, both
599
// Event and Error could be nil. Indeed, they are both nil when the event queue
600
// is empty.
601
func (c *Conn) PollForEvent() (Event, Error) {
602
	select {
603
	case everr := <-c.eventChan:
604
		return processEventOrError(everr)
605
	default:
606
		return nil, nil
607
	}
608
}
609