gitee.com/mysnapcore/mysnapd@v0.1.0/dbusutil/dbustest/dbustest.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2020 Canonical Ltd
     5   *
     6   * This program is free software: you can redistribute it and/or modify
     7   * it under the terms of the GNU General Public License version 3 as
     8   * published by the Free Software Foundation.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  package dbustest
    21  
    22  import (
    23  	"bytes"
    24  	"encoding/binary"
    25  	"fmt"
    26  	"sync/atomic"
    27  
    28  	"github.com/godbus/dbus"
    29  )
    30  
    31  // testDBusClientName is the unique name of the connected client.
    32  const testDBusClientName = ":test"
    33  
    34  // DBusHandlerFunc is the type of handler function for interacting with test DBus.
    35  //
    36  // The handler is called for each message that arrives to the bus from the test
    37  // client. The handler can respond by returning zero or more messages. Typically
    38  // one message is returned (method response or error). Additional messages can
    39  // be returned to represent signals emitted during message handling. The counter
    40  // n aids in testing a sequence of messages that is expected.
    41  //
    42  // The handler is not called for internal messages related to DBus itself.
    43  type DBusHandlerFunc func(msg *dbus.Message, n int) ([]*dbus.Message, error)
    44  
    45  // testDBusStream provides io.ReadWriteCloser for dbus.NewConn.
    46  //
    47  // Internally the stream uses separate input and output buffers.  Input buffer
    48  // is written to by DBus. Output buffer is written to internally, in response
    49  // to internal interactions with DBus or in response to message handler
    50  // function, provided by test code, returning DBus messages to send.
    51  //
    52  // Before authentication is completed buffers are used to operate a
    53  // line-oriented, text protocol. After authentication the buffers exchange DBus
    54  // messages exclusively.
    55  //
    56  // The authentication phase is handled internally and is not exposed to tests.
    57  //
    58  // Anything written to the output buffer is passed to DBus for processing.
    59  // Anything read from input buffer is decoded and either handled internally or
    60  // delegated to the handler function.
    61  //
    62  // DBus implementation we use requires Read to block while data is unavailable
    63  // so a condition variable is used to provide blocking behavior.
    64  type testDBusStream struct {
    65  	handler DBusHandlerFunc
    66  
    67  	outputBuf bytes.Buffer
    68  
    69  	output       chan []byte
    70  	closeRequest chan struct{}
    71  
    72  	closed   atomic.Value
    73  	authDone bool
    74  	n        int
    75  }
    76  
    77  func (s *testDBusStream) decodeRequest(req []byte) {
    78  	buf := bytes.NewBuffer(req)
    79  	if !s.authDone {
    80  		// Before authentication is done process the text protocol anticipating
    81  		// the TEST authentication used by NewDBusTestConn call below.
    82  		msg := buf.String()
    83  		switch msg {
    84  		case "\x00":
    85  			// initial NUL byte, ignore
    86  		case "AUTH\r\n":
    87  			s.output <- []byte("REJECTED TEST\r\n")
    88  			// XXX: this "case" can get removed once we moved to a newer version of go-dbus
    89  		case "AUTH TEST TEST\r\n":
    90  			s.output <- []byte("OK test://\r\n")
    91  		case "AUTH TEST\r\n":
    92  			s.output <- []byte("OK test://\r\n")
    93  		case "CANCEL\r\n":
    94  			s.output <- []byte("REJECTED\r\n")
    95  		case "BEGIN\r\n":
    96  			s.authDone = true
    97  		default:
    98  			panic(fmt.Errorf("unrecognized authentication message %q", msg))
    99  		}
   100  		return
   101  	}
   102  
   103  	// After authentication the buffer must contain marshaled DBus messages.
   104  	msgIn, err := dbus.DecodeMessage(buf)
   105  	if err != nil {
   106  		panic(fmt.Errorf("cannot decode incoming message: %v", err))
   107  	}
   108  	switch s.n {
   109  	case 0:
   110  		// The very first message we receive is a Hello message sent from
   111  		// NewDBusTestConn below. This message, along with the NameAcquired
   112  		// signal we send below, allow DBus to associate the client with the
   113  		// responses sent by test.
   114  		s.sendMsg(&dbus.Message{
   115  			Type: dbus.TypeMethodReply,
   116  			Headers: map[dbus.HeaderField]dbus.Variant{
   117  				dbus.FieldDestination: dbus.MakeVariant(testDBusClientName),
   118  				dbus.FieldSender:      dbus.MakeVariant("org.freedesktop.DBus"),
   119  				dbus.FieldReplySerial: dbus.MakeVariant(msgIn.Serial()),
   120  				dbus.FieldSignature:   dbus.MakeVariant(dbus.SignatureOf("")),
   121  			},
   122  			Body: []interface{}{":test"},
   123  		})
   124  		s.sendMsg(&dbus.Message{
   125  			Type: dbus.TypeSignal,
   126  			Headers: map[dbus.HeaderField]dbus.Variant{
   127  				dbus.FieldPath:        dbus.MakeVariant(dbus.ObjectPath("/org/freedesktop/DBus")),
   128  				dbus.FieldInterface:   dbus.MakeVariant("org.freedesktop.DBus"),
   129  				dbus.FieldMember:      dbus.MakeVariant("NameAcquired"),
   130  				dbus.FieldDestination: dbus.MakeVariant(testDBusClientName),
   131  				dbus.FieldSender:      dbus.MakeVariant("org.freedesktop.DBus"),
   132  				dbus.FieldSignature:   dbus.MakeVariant(dbus.SignatureOf("")),
   133  			},
   134  			Body: []interface{}{testDBusClientName},
   135  		})
   136  	default:
   137  		msgOutList, err := s.handler(msgIn, s.n-1)
   138  		if err != nil {
   139  			panic("cannot handle message: " + err.Error())
   140  		}
   141  		for _, msgOut := range msgOutList {
   142  			// Test code does not need to provide the address of the sender.
   143  			if _, ok := msgOut.Headers[dbus.FieldSender]; !ok {
   144  				msgOut.Headers[dbus.FieldSender] = dbus.MakeVariant(testDBusClientName)
   145  			}
   146  			s.sendMsg(msgOut)
   147  		}
   148  	}
   149  	s.n++
   150  }
   151  
   152  func (s *testDBusStream) sendMsg(msg *dbus.Message) {
   153  	// TODO: handle big endian if we ever get big endian machines again.
   154  	var buf bytes.Buffer
   155  	if err := msg.EncodeTo(&buf, binary.LittleEndian); err != nil {
   156  		panic(fmt.Errorf("cannot encode outgoing message: %v", err))
   157  	}
   158  	s.output <- buf.Bytes()
   159  }
   160  
   161  func (s *testDBusStream) Read(p []byte) (n int, err error) {
   162  	for {
   163  		// When the buffer is empty block until more data arrives. DBus
   164  		// continuously blocks on reading and premature empty read is treated as an
   165  		// EOF, terminating the message flow.
   166  		if s.closed.Load().(bool) {
   167  			return 0, fmt.Errorf("stream is closed")
   168  		}
   169  		if s.outputBuf.Len() > 0 {
   170  			return s.outputBuf.Read(p)
   171  		}
   172  		select {
   173  		case data := <-s.output:
   174  			// just accumulate the data in the output buffer
   175  			s.outputBuf.Write(data)
   176  		case <-s.closeRequest:
   177  			s.closed.Store(true)
   178  		}
   179  	}
   180  }
   181  
   182  func (s *testDBusStream) Write(p []byte) (n int, err error) {
   183  	for {
   184  		select {
   185  		case <-s.closeRequest:
   186  			s.closed.Store(true)
   187  		default:
   188  			if s.closed.Load().(bool) {
   189  				return 0, fmt.Errorf("stream is closed")
   190  			}
   191  			s.decodeRequest(p)
   192  			return len(p), nil
   193  		}
   194  	}
   195  }
   196  
   197  func (s *testDBusStream) Close() error {
   198  	s.closeRequest <- struct{}{}
   199  	return nil
   200  }
   201  
   202  func (s *testDBusStream) InjectMessage(msg *dbus.Message) {
   203  	s.sendMsg(msg)
   204  }
   205  
   206  func newTestDBusStream(handler DBusHandlerFunc) *testDBusStream {
   207  	s := &testDBusStream{
   208  		handler:      handler,
   209  		output:       make(chan []byte, 1),
   210  		closeRequest: make(chan struct{}, 1),
   211  	}
   212  	s.closed.Store(false)
   213  	return s
   214  }
   215  
   216  // testAuth implements DBus authentication protocol used during testing.
   217  type testAuth struct{}
   218  
   219  func (a *testAuth) FirstData() (name, resp []byte, status dbus.AuthStatus) {
   220  	return []byte("TEST"), []byte("TEST"), dbus.AuthOk
   221  }
   222  
   223  func (a *testAuth) HandleData(data []byte) (resp []byte, status dbus.AuthStatus) {
   224  	return []byte(""), dbus.AuthOk
   225  }
   226  
   227  type InjectMessageFunc func(msg *dbus.Message)
   228  
   229  // InjectableConnection returns a DBus connection for writing unit tests and a
   230  // function that can be used to inject messages that will be received by the
   231  // test client.
   232  func InjectableConnection(handler DBusHandlerFunc) (*dbus.Conn, InjectMessageFunc, error) {
   233  	testDBusStream := newTestDBusStream(handler)
   234  	conn, err := dbus.NewConn(testDBusStream)
   235  	if err != nil {
   236  		return nil, nil, err
   237  	}
   238  	if err = conn.Auth([]dbus.Auth{&testAuth{}}); err != nil {
   239  		_ = conn.Close()
   240  		return nil, nil, err
   241  	}
   242  	if err = conn.Hello(); err != nil {
   243  		_ = conn.Close()
   244  		return nil, nil, err
   245  	}
   246  	return conn, testDBusStream.InjectMessage, nil
   247  }
   248  
   249  // Connection returns a DBus connection for writing unit tests.
   250  //
   251  // The handler function is called for each message sent to the bus. It can
   252  // return any number of messages to send in response.
   253  func Connection(handler DBusHandlerFunc) (*dbus.Conn, error) {
   254  	conn, _, err := InjectableConnection(handler)
   255  	return conn, err
   256  }