1 /*
2  * Copyright (C) 2020 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.testutils
18 
19 import android.Manifest.permission.MANAGE_TEST_NETWORKS
20 import android.net.TestNetworkInterface
21 import android.net.TestNetworkManager
22 import android.os.Handler
23 import android.os.HandlerThread
24 import androidx.test.platform.app.InstrumentationRegistry
25 import org.junit.rules.TestRule
26 import org.junit.runner.Description
27 import org.junit.runners.model.Statement
28 import kotlin.test.assertFalse
29 import kotlin.test.fail
30 
31 private const val HANDLER_TIMEOUT_MS = 10_000L
32 
33 /**
34  * A [TestRule] that sets up a [TapPacketReader] on a [TestNetworkInterface] for use in the test.
35  *
36  * @param maxPacketSize Maximum size of packets read in the [TapPacketReader] buffer.
37  * @param autoStart Whether to initialize the interface and start the reader automatically for every
38  *                  test. If false, each test must either call start() and stop(), or be annotated
39  *                  with TapPacketReaderTest before using the reader or interface.
40  */
41 class TapPacketReaderRule @JvmOverloads constructor(
42     private val maxPacketSize: Int = 1500,
43     private val autoStart: Boolean = true
44 ) : TestRule {
45     // Use lateinit as the below members can't be initialized in the rule constructor (the
46     // InstrumentationRegistry may not be ready), but from the point of view of test cases using
47     // this rule with autoStart = true, the members are always initialized (in setup/test/teardown):
48     // tests cases should be able use them directly.
49     // lateinit also allows getting good exceptions detailing what went wrong if the members are
50     // referenced before they could be initialized (typically if autoStart is false and the test
51     // does not call start or use @TapPacketReaderTest).
52     lateinit var iface: TestNetworkInterface
53     lateinit var reader: TapPacketReader
54 
55     @Volatile
56     private var readerRunning = false
57 
58     /**
59      * Indicates that the [TapPacketReaderRule] should initialize its [TestNetworkInterface] and
60      * start the [TapPacketReader] before the test, and tear them down afterwards.
61      *
62      * For use when [TapPacketReaderRule] is created with autoStart = false.
63      */
64     annotation class TapPacketReaderTest
65 
66     /**
67      * Initialize the tap interface and start the [TapPacketReader].
68      *
69      * Tests using this method must also call [stop] before exiting.
70      * @param handler Handler to run the reader on. Callers are responsible for safely terminating
71      *                the handler when the test ends. If null, a handler thread managed by the
72      *                rule will be used.
73      */
74     @JvmOverloads
75     fun start(handler: Handler? = null) {
76         if (this::iface.isInitialized) {
77             fail("${TapPacketReaderRule::class.java.simpleName} was already started")
78         }
79 
80         val ctx = InstrumentationRegistry.getInstrumentation().context
81         iface = runAsShell(MANAGE_TEST_NETWORKS) {
82             val tnm = ctx.getSystemService(TestNetworkManager::class.java)
83                     ?: fail("Could not obtain the TestNetworkManager")
84             tnm.createTapInterface()
85         }
86         val usedHandler = handler ?: HandlerThread(
87                 TapPacketReaderRule::class.java.simpleName).apply { start() }.threadHandler
88         reader = TapPacketReader(usedHandler, iface.fileDescriptor.fileDescriptor, maxPacketSize)
89         reader.startAsyncForTest()
90         readerRunning = true
91     }
92 
93     /**
94      * Stop the [TapPacketReader].
95      *
96      * Tests calling [start] must call this method before exiting. If a handler was specified in
97      * [start], all messages on that handler must also be processed after calling this method and
98      * before exiting.
99      *
100      * If [start] was not called, calling this method is a no-op.
101      */
102     fun stop() {
103         // The reader may not be initialized if the test case did not use the rule, even though
104         // other test cases in the same class may be using it (so test classes may call stop in
105         // tearDown even if start is not called for all test cases).
106         if (!this::reader.isInitialized) return
107         reader.handler.post {
108             reader.stop()
109             readerRunning = false
110         }
111     }
112 
113     override fun apply(base: Statement, description: Description): Statement {
114         return TapReaderStatement(base, description)
115     }
116 
117     private inner class TapReaderStatement(
118         private val base: Statement,
119         private val description: Description
120     ) : Statement() {
121         override fun evaluate() {
122             val shouldStart = autoStart ||
123                     description.getAnnotation(TapPacketReaderTest::class.java) != null
124             if (shouldStart) {
125                 start()
126             }
127 
128             try {
129                 base.evaluate()
130             } finally {
131                 if (shouldStart) {
132                     stop()
133                     reader.handler.looper.apply {
134                         quitSafely()
135                         thread.join(HANDLER_TIMEOUT_MS)
136                         assertFalse(thread.isAlive,
137                                 "HandlerThread did not exit within $HANDLER_TIMEOUT_MS ms")
138                     }
139                 }
140 
141                 if (this@TapPacketReaderRule::iface.isInitialized) {
142                     iface.fileDescriptor.close()
143                 }
144             }
145 
146             assertFalse(readerRunning,
147                     "stop() was not called, or the provided handler did not process the stop " +
148                     "message before the test ended. If not using autostart, make sure to call " +
149                     "stop() after the test. If a handler is specified in start(), make sure all " +
150                     "messages are processed after calling stop(), before quitting (for example " +
151                     "by using HandlerThread#quitSafely and HandlerThread#join).")
152         }
153     }
154 }