1 /*
2  * Copyright (C) 2023 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 android.util;
18 
19 import static android.util.StatsEvent.TYPE_ATTRIBUTION_CHAIN;
20 import static android.util.StatsEvent.TYPE_BOOLEAN;
21 import static android.util.StatsEvent.TYPE_BYTE_ARRAY;
22 import static android.util.StatsEvent.TYPE_ERRORS;
23 import static android.util.StatsEvent.TYPE_FLOAT;
24 import static android.util.StatsEvent.TYPE_INT;
25 import static android.util.StatsEvent.TYPE_LIST;
26 import static android.util.StatsEvent.TYPE_LONG;
27 import static android.util.StatsEvent.TYPE_STRING;
28 import static android.util.proto.ProtoOutputStream.FIELD_COUNT_REPEATED;
29 import static android.util.proto.ProtoOutputStream.FIELD_COUNT_SINGLE;
30 import static android.util.proto.ProtoOutputStream.FIELD_TYPE_FLOAT;
31 import static android.util.proto.ProtoOutputStream.FIELD_TYPE_INT32;
32 import static android.util.proto.ProtoOutputStream.FIELD_TYPE_INT64;
33 import static android.util.proto.ProtoOutputStream.FIELD_TYPE_MESSAGE;
34 import static android.util.proto.ProtoOutputStream.FIELD_TYPE_STRING;
35 import static java.nio.charset.StandardCharsets.UTF_8;
36 
37 import android.util.proto.ProtoOutputStream;
38 import com.android.os.AtomsProto.Atom;
39 import com.google.protobuf.InvalidProtocolBufferException;
40 import java.nio.ByteBuffer;
41 import java.nio.ByteOrder;
42 
43 public final class StatsEventTestUtils {
44     private static final int ATTRIBUTION_UID_FIELD = 1;
45     private static final int ATTRIBUTION_TAG_FIELD = 2;
46 
StatsEventTestUtils()47     private StatsEventTestUtils() {
48     } // no instances.
49 
50     // Convert StatsEvent to MessageLite representation of Atom.
51     // Calls StatsEvent#release; No further actions should be taken on the StatsEvent
52     // object.
convertToAtom(StatsEvent statsEvent)53     public static Atom convertToAtom(StatsEvent statsEvent) throws InvalidProtocolBufferException {
54         return Atom.parseFrom(getProtoBytes(statsEvent));
55     }
56 
57     // Convert StatsEvent to serialized proto representation of Atom.
58     // Calls StatsEvent#release; No further actions should be taken on the StatsEvent
59     // object.
getProtoBytes(StatsEvent statsEvent)60     private static byte[] getProtoBytes(StatsEvent statsEvent) {
61         try {
62             ByteBuffer buf = ByteBuffer.wrap(statsEvent.getBytes()).order(ByteOrder.LITTLE_ENDIAN);
63             buf.get(); // Payload starts with TYPE_OBJECT.
64 
65             // Read number of elements at the root level.
66             byte fieldsRemaining = buf.get();
67             if (fieldsRemaining < 2) {
68                 // Each StatsEvent should at least have a timestamp and atom ID.
69                 throw new IllegalArgumentException("StatsEvent should have more than 2 elements.");
70             }
71 
72             // Read timestamp.
73             if (buf.get() != TYPE_LONG) {
74                 // Timestamp should be TYPE_LONG
75                 throw new IllegalArgumentException("StatsEvent does not have timestamp.");
76             }
77             buf.getLong(); // Read elapsed timestamp.
78             fieldsRemaining--;
79 
80             // Read atom ID.
81             FieldMetadata fieldMetadata = parseFieldMetadata(buf);
82             if (fieldMetadata.typeId != TYPE_INT) {
83                 // atom ID should be an integer.
84                 throw new IllegalArgumentException("StatsEvent does not have an atom ID.");
85             }
86             int atomId = buf.getInt();
87             skipAnnotations(buf, fieldMetadata.annotationCount);
88             fieldsRemaining--;
89 
90             ProtoOutputStream proto = new ProtoOutputStream();
91             long atomToken = proto.start(FIELD_TYPE_MESSAGE | FIELD_COUNT_SINGLE | atomId);
92 
93             // Read atom fields.
94             for (int tag = 1; tag <= fieldsRemaining; tag++) {
95                 fieldMetadata = parseFieldMetadata(buf);
96                 parseField(fieldMetadata.typeId, FIELD_COUNT_SINGLE, tag, buf, proto);
97                 skipAnnotations(buf, fieldMetadata.annotationCount);
98             }
99 
100             // We should have parsed all bytes in StatsEvent at this point.
101             if (buf.position() != statsEvent.getNumBytes()) {
102                 throw new IllegalArgumentException("Unexpected bytes in StatsEvent");
103             }
104 
105             proto.end(atomToken);
106             return proto.getBytes();
107         } finally {
108             statsEvent.release();
109         }
110     }
111 
parseField( byte typeId, long fieldCount, int tag, ByteBuffer buf, ProtoOutputStream proto)112     private static void parseField(
113             byte typeId, long fieldCount, int tag, ByteBuffer buf, ProtoOutputStream proto) {
114         switch (typeId) {
115             case TYPE_INT:
116                 proto.write(FIELD_TYPE_INT32 | fieldCount | tag, buf.getInt());
117                 break;
118             case TYPE_LONG:
119                 proto.write(FIELD_TYPE_INT64 | fieldCount | tag, buf.getLong());
120                 break;
121             case TYPE_STRING:
122                 String value = new String(getByteArrayFromByteBuffer(buf), UTF_8);
123                 proto.write(FIELD_TYPE_STRING | fieldCount | tag, value);
124                 break;
125             case TYPE_FLOAT:
126                 proto.write(FIELD_TYPE_FLOAT | fieldCount | tag, buf.getFloat());
127                 break;
128             case TYPE_BOOLEAN:
129                 proto.write(FIELD_TYPE_INT32 | fieldCount | tag, buf.get());
130                 break;
131             case TYPE_ATTRIBUTION_CHAIN:
132                 byte numNodes = buf.get();
133                 for (byte i = 1; i <= numNodes; i++) {
134                     long token = proto.start(FIELD_TYPE_MESSAGE | FIELD_COUNT_REPEATED | tag);
135                     proto.write(FIELD_TYPE_INT32 | FIELD_COUNT_SINGLE | ATTRIBUTION_UID_FIELD,
136                             buf.getInt());
137                     String tagName = new String(getByteArrayFromByteBuffer(buf), UTF_8);
138                     proto.write(FIELD_TYPE_STRING | FIELD_COUNT_SINGLE | ATTRIBUTION_TAG_FIELD,
139                             tagName);
140                     proto.end(token);
141                 }
142                 break;
143             case TYPE_BYTE_ARRAY:
144                 byte[] byteArray = getByteArrayFromByteBuffer(buf);
145                 proto.write(FIELD_TYPE_MESSAGE | FIELD_COUNT_SINGLE | tag, byteArray);
146                 break;
147             case TYPE_LIST:
148                 byte numItems = buf.get();
149                 byte listTypeId = buf.get();
150                 for (byte i = 1; i <= numItems; i++) {
151                     parseField(listTypeId, FIELD_COUNT_REPEATED, tag, buf, proto);
152                 }
153                 break;
154             case TYPE_ERRORS:
155                 int errorMask = buf.getInt();
156                 throw new IllegalArgumentException("StatsEvent has error(s): " + errorMask);
157             default:
158                 throw new IllegalArgumentException(
159                         "Invalid typeId encountered while parsing StatsEvent: " + typeId);
160         }
161     }
162 
getByteArrayFromByteBuffer(ByteBuffer buf)163     private static byte[] getByteArrayFromByteBuffer(ByteBuffer buf) {
164         final int numBytes = buf.getInt();
165         byte[] bytes = new byte[numBytes];
166         buf.get(bytes);
167         return bytes;
168     }
169 
skipAnnotations(ByteBuffer buf, int annotationCount)170     private static void skipAnnotations(ByteBuffer buf, int annotationCount) {
171         for (int i = 1; i <= annotationCount; i++) {
172             buf.get(); // read annotation ID.
173             byte annotationType = buf.get();
174             if (annotationType == TYPE_INT) {
175                 buf.getInt(); // read and drop int annotation value.
176             } else if (annotationType == TYPE_BOOLEAN) {
177                 buf.get(); // read and drop byte annotation value.
178             } else {
179                 throw new IllegalArgumentException("StatsEvent has an invalid annotation.");
180             }
181         }
182     }
183 
parseFieldMetadata(ByteBuffer buf)184     private static FieldMetadata parseFieldMetadata(ByteBuffer buf) {
185         FieldMetadata fieldMetadata = new FieldMetadata();
186         fieldMetadata.typeId = buf.get();
187         fieldMetadata.annotationCount = (byte) (fieldMetadata.typeId >> 4);
188         fieldMetadata.typeId &= (byte) 0x0F;
189 
190         return fieldMetadata;
191     }
192 
193     private static class FieldMetadata {
194         byte typeId;
195         byte annotationCount;
196     }
197 }
198